458 lines
18 KiB
Python
458 lines
18 KiB
Python
# Copyright (c) 2018 Orange. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"); you may
|
|
# not use this file except in compliance with the License. You may obtain
|
|
# a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
# License for the specific language governing permissions and limitations
|
|
# under the License.
|
|
#
|
|
|
|
# Conflict with flake8
|
|
# pylint: disable = bad-continuation
|
|
|
|
"""A theory that contains rules that must be treated by Z3."""
|
|
|
|
import time
|
|
|
|
import logging
|
|
import six
|
|
|
|
from congress.datalog import base
|
|
from congress.datalog import compile as ast
|
|
from congress.datalog import nonrecursive
|
|
from congress.datalog import ruleset
|
|
from congress.datalog import unify
|
|
from congress import exception
|
|
from congress.z3 import typechecker
|
|
from congress.z3 import z3builtins
|
|
from congress.z3 import z3types
|
|
|
|
# pylint: disable = ungrouped-imports
|
|
MYPY = False
|
|
if MYPY:
|
|
# pylint: disable = unused-import
|
|
from congress.datalog import topdown # noqa
|
|
from mypy_extensions import TypedDict # noqa
|
|
from typing import Dict, Callable, Optional, Union, List, Any, Tuple # noqa
|
|
import z3 # noqa
|
|
Z3_RESULT = Tuple[Union[bool, List[List[z3.ExprRef]]],
|
|
List[ast.Variable],
|
|
List[z3types.Z3Type]]
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
Z3_ENGINE_OPTIONS = {'engine': 'datalog'}
|
|
|
|
Z3OPT = z3types.z3
|
|
|
|
INTER_COMPILE_DELAY = 60.0
|
|
|
|
|
|
def cycle_not_contained_in_z3(theories, cycles):
|
|
# type: (Dict[str, base.Theory], List[List[str]]) -> bool
|
|
"""Check that there is a true cycle not through Z3 theory
|
|
|
|
A cycle is irreducible if it contains at least one element which is not a
|
|
Z3Theory for which recursion is allowed. Cycles are represented by lists of
|
|
qualified table names.
|
|
"""
|
|
acceptables = [
|
|
th.name
|
|
for th in six.itervalues(theories)
|
|
if isinstance(th, Z3Theory)]
|
|
return any(fullname[:fullname.index(':')] not in acceptables
|
|
for cycle in cycles for fullname in cycle)
|
|
|
|
|
|
# TODO(pcregut): Object constants should evolve to use the type system
|
|
# rather than custom types.
|
|
def congress_constant(val):
|
|
"""Creates an object constant from a value using its type"""
|
|
if isinstance(val, six.string_types):
|
|
typ = ast.ObjectConstant.STRING
|
|
elif isinstance(val, int):
|
|
typ = ast.ObjectConstant.INTEGER
|
|
elif isinstance(val, float):
|
|
typ = ast.ObjectConstant.FLOAT
|
|
else:
|
|
val = str(val)
|
|
typ = ast.ObjectConstant.STRING
|
|
return ast.ObjectConstant(val, typ)
|
|
|
|
|
|
def retrieve(theory, tablename):
|
|
# type: (topdown.TopDownTheory, str) -> List[ast.Literal]
|
|
"""Retrieves all the values of an external table.
|
|
|
|
Performs a select on the theory with a query computed from the schema
|
|
of the table.
|
|
"""
|
|
arity = theory.schema.arity(tablename)
|
|
table = ast.Tablename(tablename, theory.name)
|
|
args = [ast.Variable('X' + str(i)) for i in range(arity)]
|
|
query = ast.Literal(table, args)
|
|
return theory.select(query)
|
|
|
|
|
|
class Z3Theory(nonrecursive.RuleHandlingMixin, base.Theory):
|
|
"""Theory for Z3 engine
|
|
|
|
Z3Theory is a datalog theory interpreted by the Z3 engine instead of
|
|
the usual congress internal engine.
|
|
"""
|
|
|
|
def __init__(self, name=None, abbr=None,
|
|
schema=None, theories=None, desc=None, owner=None):
|
|
super(Z3Theory, self).__init__(
|
|
name=name, abbr=abbr, theories=theories,
|
|
schema=ast.Schema() if schema is None else schema,
|
|
desc=desc, owner=owner)
|
|
LOG.info('z3theory: create %s', name)
|
|
self.kind = base.Z3_POLICY_TYPE
|
|
self.rules = ruleset.RuleSet()
|
|
self.dirty = False
|
|
self.z3context = None
|
|
Z3Context.get_context().register(self)
|
|
|
|
def select(self, query, find_all=True):
|
|
"""Performs a query"""
|
|
return self.z3context.select(self, query, find_all)
|
|
|
|
def arity(self, tablename, modal=None):
|
|
"""Arity of a table"""
|
|
return self.schema.arity(tablename)
|
|
|
|
def drop(self):
|
|
"""To call when the theory is forgotten"""
|
|
self.z3context.drop(self)
|
|
|
|
def _top_down_eval(self,
|
|
context, # type: topdown.TopDownTheory.TopDownContext
|
|
caller # type: topdown.TopDownTheory.TopDownCaller
|
|
):
|
|
# type: (...) -> bool
|
|
"""Evaluation entry point for the non recursive engine
|
|
|
|
We must compute unifiers and clear off as soon as we can
|
|
giving back control to the theory context.
|
|
Returns true if we only need one binding and it has been found,
|
|
false otherwise.
|
|
"""
|
|
raw_lit = context.literals[context.literal_index]
|
|
query_lit = raw_lit.plug(context.binding)
|
|
answers, bvars, translators = self.z3context.eval(self, query_lit)
|
|
if isinstance(answers, bool):
|
|
if answers:
|
|
return (context.theory._top_down_finish(context, caller)
|
|
and not caller.find_all)
|
|
return False
|
|
for answer in answers:
|
|
changes = []
|
|
for (val, var, trans) in six.moves.zip(answer, bvars, translators):
|
|
chg = context.binding.add(var, trans.to_os(val), None)
|
|
changes.append(chg)
|
|
context.theory._top_down_finish(context, caller)
|
|
unify.undo_all(changes)
|
|
if not caller.find_all:
|
|
return True
|
|
return False
|
|
|
|
|
|
class Z3Context(object):
|
|
"""An instance of Z3 defined first by its execution context"""
|
|
|
|
_singleton = None
|
|
|
|
def __init__(self):
|
|
self.context = Z3OPT.Fixedpoint()
|
|
self.context.set(**Z3_ENGINE_OPTIONS)
|
|
self.z3theories = {} # type: Dict[str, Z3Theory]
|
|
self.relations = {} # type: Dict[str, z3.Function]
|
|
# back pointer on all theories extracted from registered theory.
|
|
self.theories = None # type: Dict[str, topdown.TopDownTheory]
|
|
self.externals = set() # type: Set[Tuple[str, str]]
|
|
self.type_registry = z3types.TypeRegistry()
|
|
self.last_compiled = 0
|
|
|
|
def register(self, theory):
|
|
# type: (Z3Theory) -> None
|
|
"""Registers a Z3 theory in the context"""
|
|
if self.theories is None:
|
|
self.theories = theory.theories
|
|
theory.z3context = self
|
|
self.z3theories[theory.name] = theory
|
|
|
|
def drop(self, theory):
|
|
# type: (Z3Theory) -> None
|
|
"""Unregister a Z3 theory from the context"""
|
|
del self.z3theories[theory.name]
|
|
|
|
@staticmethod
|
|
def get_context():
|
|
# type: () -> Z3Context
|
|
"""Gives back the unique instance of this class.
|
|
|
|
Users should not use the class constructor but this method.
|
|
"""
|
|
if Z3Context._singleton is None:
|
|
Z3Context._singleton = Z3Context()
|
|
return Z3Context._singleton
|
|
|
|
def eval(self,
|
|
theory, # type: Z3Theory
|
|
query # type: ast.Literal
|
|
):
|
|
# type: (...) -> Z3_RESULT
|
|
"""Solves a query and gives back a raw result
|
|
|
|
Result is in Z3 ast format with a translator
|
|
"""
|
|
theories_changed = any(t.dirty for t in self.z3theories.values())
|
|
# TODO(pcregut): replace either with an option or find something
|
|
# better for the refresh of datasources.
|
|
needs_refresh = time.time() - self.last_compiled > INTER_COMPILE_DELAY
|
|
if theories_changed or needs_refresh:
|
|
# There is no reset on Z3 context. Replace with a new one.
|
|
self.context = Z3OPT.Fixedpoint()
|
|
self.context.set(**Z3_ENGINE_OPTIONS)
|
|
type_env = self.typecheck()
|
|
self.compile_all(type_env)
|
|
self.synchronize_external()
|
|
z3query = self.compile_query(theory, query)
|
|
self.context.query(z3query)
|
|
z3answer = self.context.get_answer()
|
|
answer = z3types.z3_to_array(z3answer)
|
|
typ_args = theory.schema.types(query.table.table)
|
|
variables = [] # type: List[ast.Variable]
|
|
translators = [] # type: List[z3types.Z3Type]
|
|
for arg, typ_arg in six.moves.zip(query.arguments, typ_args):
|
|
if isinstance(arg, ast.Variable) and arg not in variables:
|
|
translators.append(
|
|
self.type_registry.get_translator(str(typ_arg.type)))
|
|
variables.append(arg)
|
|
return (answer, variables, translators)
|
|
|
|
def select(self, theory, query, find_all):
|
|
# type: (Z3Theory, ast.Literal, bool) -> List[ast.Literal]
|
|
"""Query a theory"""
|
|
(answer, variables, trans) = self.eval(theory, query)
|
|
pattern = [
|
|
variables.index(arg) if isinstance(arg, ast.Variable) else arg
|
|
for arg in query.arguments]
|
|
|
|
def plug(row):
|
|
"""Plugs in found values in query litteral"""
|
|
args = [
|
|
(congress_constant(trans[arg].to_os(row[arg]))
|
|
if isinstance(arg, int) else arg)
|
|
for arg in pattern]
|
|
return ast.Literal(query.table, args)
|
|
|
|
if isinstance(answer, bool):
|
|
return [query] if answer else []
|
|
if find_all:
|
|
result = [plug(row) for row in answer]
|
|
else:
|
|
result = [plug(answer[0])]
|
|
return result
|
|
|
|
def declare_table(self, theory, tablename):
|
|
"""Declares a new table in Z3 context"""
|
|
fullname = theory.name + ':' + tablename
|
|
if fullname in self.relations:
|
|
return
|
|
typ_args = theory.schema.types(tablename)
|
|
param_types = [
|
|
self.type_registry.get_type(str(tArg.type))
|
|
for tArg in typ_args]
|
|
param_types.append(Z3OPT.BoolSort())
|
|
relation = Z3OPT.Function(fullname, *param_types)
|
|
self.context.register_relation(relation)
|
|
self.relations[fullname] = relation
|
|
|
|
def declare_tables(self):
|
|
"""Declares all tables defined in Z3 context"""
|
|
for theory in six.itervalues(self.z3theories):
|
|
for tablename in theory.schema.map.keys():
|
|
self.declare_table(theory, tablename)
|
|
|
|
def declare_external_tables(self):
|
|
"""Declares tables from other theories used in Z3 context"""
|
|
def declare_for_lit(lit):
|
|
"""Declares the table of a litteral if necessary"""
|
|
service = lit.table.service
|
|
table = lit.table.table
|
|
if (service is not None and service != 'builtin' and
|
|
service not in self.z3theories):
|
|
self.externals.add((service, table))
|
|
for theory in six.itervalues(self.z3theories):
|
|
for rules in six.itervalues(theory.rules.rules):
|
|
for rule in rules:
|
|
for lit in rule.body:
|
|
declare_for_lit(lit)
|
|
for (service, table) in self.externals:
|
|
self.declare_table(self.theories[service], table)
|
|
|
|
def compile_facts(self, theory):
|
|
# type: (Z3Theory) -> None
|
|
"""Compiles the facts of a theory in Z3 context"""
|
|
for tname, facts in six.iteritems(theory.rules.facts):
|
|
translators = [
|
|
self.type_registry.get_translator(str(arg_type.type))
|
|
for arg_type in theory.schema.types(tname)]
|
|
fullname = theory.name + ':' + tname
|
|
z3func = self.relations[fullname]
|
|
for fact in facts:
|
|
z3args = (tr.to_z3(v, strict=True)
|
|
for (v, tr) in six.moves.zip(fact, translators))
|
|
z3fact = z3func(*z3args)
|
|
self.context.fact(z3fact)
|
|
|
|
def compile_atoms(self,
|
|
type_env,
|
|
theory, # type: Z3Theory
|
|
head, # type: ast.Literal
|
|
body # type: List[ast.Literal]
|
|
):
|
|
# type: (...) -> Tuple[z3.Const, z3.ExprRef, List[z3.ExprRef]]
|
|
"""Compile a list of atoms belonging to a single variable scope
|
|
|
|
As it is used mainly for rules, the head is distinguished.
|
|
"""
|
|
variables = {} # type: Dict[str, z3.Const]
|
|
z3vars = []
|
|
|
|
def compile_expr(expr, translator):
|
|
"""Compiles an expression to Z3"""
|
|
if isinstance(expr, ast.Variable):
|
|
name = expr.name
|
|
if name in variables:
|
|
return variables[name]
|
|
var = Z3OPT.Const(name, translator.type())
|
|
variables[name] = var
|
|
z3vars.append(var)
|
|
return var
|
|
elif isinstance(expr, ast.ObjectConstant):
|
|
return translator.to_z3(expr.name)
|
|
else:
|
|
raise exception.PolicyException(
|
|
"Expr {} not handled by Z3".format(expr))
|
|
|
|
def compile_atom(literal, pos=-1):
|
|
"""Compiles an atom in Z3"""
|
|
name = literal.table.table
|
|
svc = literal.table.service
|
|
if svc == 'builtin':
|
|
translators = [
|
|
self.type_registry.get_translator(str(arg_type['type']))
|
|
for arg_type in type_env[pos]
|
|
]
|
|
fullname = 'builtin:'+name
|
|
else:
|
|
lit_theory = theory if svc is None else self.theories[svc]
|
|
translators = [
|
|
self.type_registry.get_translator(str(arg_type.type))
|
|
for arg_type in lit_theory.schema.types(name)]
|
|
fullname = lit_theory.name + ":" + name
|
|
try:
|
|
z3func = (
|
|
z3builtins.BUILTINS[name].z3 if svc == 'builtin'
|
|
else self.relations[fullname])
|
|
z3args = (compile_expr(arg, tr)
|
|
for (arg, tr) in six.moves.zip(literal.arguments,
|
|
translators))
|
|
z3lit = z3func(*z3args)
|
|
return (Z3OPT.Not(z3lit) if literal.negated
|
|
else z3lit)
|
|
except KeyError:
|
|
raise exception.PolicyException(
|
|
"Z3: Relation %s not registered" % fullname)
|
|
|
|
z3head = compile_atom(head)
|
|
z3body = [compile_atom(atom, pos) for (pos, atom) in enumerate(body)]
|
|
# We give back variables explicitely and do not rely on declare_var and
|
|
# abstract. Otherwise rules are cluttered with useless variables.
|
|
return (z3vars, z3head, z3body)
|
|
|
|
def compile_rule(self, type_env, theory, rule):
|
|
# type: (typechecker.GEN_TYPE_ENV, Z3Theory, ast.Rule) -> None
|
|
"""compiles a single rule
|
|
|
|
:param theory: the theory containing the rule
|
|
:param rule: the rule to compile.
|
|
"""
|
|
z3vars, z3head, z3body = self.compile_atoms(
|
|
type_env.get(rule.id, {}), theory, rule.head, rule.body)
|
|
term1 = (z3head if z3body == []
|
|
else Z3OPT.Implies(Z3OPT.And(*z3body), z3head))
|
|
term2 = term1 if z3vars == [] else Z3OPT.ForAll(z3vars, term1)
|
|
self.context.rule(term2)
|
|
|
|
def compile_query(self, theory, literal):
|
|
# type: (Z3Theory, ast.Literal) -> z3.ExprRef
|
|
"""compiles a query litteral
|
|
|
|
:param theory: theory used as the context of the query
|
|
:param litteral: the query
|
|
:returns: an existentially quantified litteral in Z3 format.
|
|
"""
|
|
z3vars, z3head, _ = self.compile_atoms({}, theory, literal, [])
|
|
return z3head if z3vars == [] else Z3OPT.Exists(z3vars, z3head)
|
|
|
|
def compile_theory(self, type_env, theory):
|
|
# type: (typechecker.GEN_TYPE_ENV, Z3Theory) -> None
|
|
"""Compiles all the rules of a theory
|
|
|
|
:param theory: theory to compile. Will be marked clean after.
|
|
"""
|
|
self.compile_facts(theory)
|
|
for rules in six.itervalues(theory.rules.rules):
|
|
for rule in rules:
|
|
self.compile_rule(type_env, theory, rule)
|
|
theory.dirty = False
|
|
|
|
def compile_all(self, type_env):
|
|
"""Compile all Z3 theories"""
|
|
self.relations = {}
|
|
self.externals.clear()
|
|
self.declare_tables()
|
|
self.declare_external_tables()
|
|
for theory in six.itervalues(self.z3theories):
|
|
self.compile_theory(type_env, theory)
|
|
self.last_compiled = time.time()
|
|
|
|
def typecheck(self):
|
|
"""Typechecker for rules defined"""
|
|
typer = typechecker.Typechecker(
|
|
self.z3theories.values(), self.theories)
|
|
return typer.type_all()
|
|
|
|
def inject(self, theoryname, tablename):
|
|
# type: (str, str) -> None
|
|
"""Inject the values of an external table in the Z3 Context.
|
|
|
|
Loops over the literal retrieved from a standard query.
|
|
"""
|
|
theory = self.theories[theoryname]
|
|
translators = [
|
|
self.type_registry.get_translator(str(arg_type.type))
|
|
for arg_type in theory.schema.types(tablename)]
|
|
fullname = theory.name + ':' + tablename
|
|
z3func = self.relations[fullname]
|
|
for lit in retrieve(theory, tablename):
|
|
z3args = (tr.to_z3(v.name, strict=True)
|
|
for (v, tr) in six.moves.zip(lit.arguments, translators))
|
|
z3fact = z3func(*z3args)
|
|
self.context.fact(z3fact)
|
|
|
|
def synchronize_external(self):
|
|
"""Synchronize all external tables"""
|
|
for (theoryname, tablename) in self.externals:
|
|
self.inject(theoryname, tablename)
|