Files

692 lines
22 KiB
Python

#!/usr/bin/env python3
"""datalog.py -- a self-contained Datalog engine.
Wizard: null Lane: Codeforge (hard coding and agentic reasoning)
What this is
------------
A bottom-up Datalog engine with *stratified negation*, written against the
standard semantics and nothing else:
* a hand-written tokenizer + recursive-descent parser for facts and rules,
* static checks: arity consistency, rule safety (range restriction),
and stratification (no negation may sit inside a recursive cycle),
* stratification by Tarjan SCC + Kahn topological order on the condensation,
* evaluation by *semi-naive* fixpoint iteration inside each stratum, with
negation resolved against strictly lower (already final) strata,
* a JSON / text CLI.
Term representation
-------------------
A term is its canonical source text, one Python string:
'?X' variable X (leading '?' marks a variable)
'?_' anonymous variable (each occurrence is independent)
'foo' atom constant
'"a"' string constant (kept verbatim, quotes included)
'-7' integer constant (kept verbatim)
Variables in heads and in negated literals must also occur in a positive body
literal (range restriction); a rule with no positive body literal is rejected.
Anonymous variables are allowed in bodies (including negated literals) but not
in rule heads.
Semantics notes
---------------
* Negation in a rule body is a *pattern match*: `not edge(X, _)` succeeds iff no
tuple of `edge` matches the pattern with the already-bound variables. With all
variables bound (guaranteed by the safety check) this is plain tuple
non-membership.
* Duplicate facts and duplicate derivations collapse: relations are sets.
* The result is independent of the order of facts and rules.
Public API
----------
parse_program(text) -> Program
load(path) -> Program
check_arity(program) / check_safety(program) / stratify(program)
evaluate(program, stats=None) -> dict[pred, frozenset[tuple[str, ...]]]
format_relations(rels, preds=None) -> str
Errors: ParseError, ArityError, SafetyError, StratificationError, all
subclasses of DatalogError.
Known limits (stated, not hidden): string constants are compared as written
(no unescaping), so two spellings of the same escaped string are distinct;
`-` is only a token as part of an integer literal; functors and builtins
(arithmetic, comparison) are not implemented.
"""
from __future__ import annotations
import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field
__all__ = [
"DatalogError",
"ParseError",
"ArityError",
"SafetyError",
"StratificationError",
"Literal",
"Rule",
"Program",
"parse_program",
"load",
"check_arity",
"check_safety",
"stratify",
"evaluate",
"format_relations",
"main",
]
class DatalogError(Exception):
"""Base class for every error this engine raises."""
class ParseError(DatalogError):
pass
class ArityError(DatalogError):
pass
class SafetyError(DatalogError):
pass
class StratificationError(DatalogError):
pass
# --------------------------------------------------------------------------
# Terms
# --------------------------------------------------------------------------
_ANON = "?_"
def is_var(term: str) -> bool:
return term.startswith("?")
def is_anon(term: str) -> bool:
return term == _ANON
def var_name(term: str) -> str:
return term[1:]
# --------------------------------------------------------------------------
# Syntax tree
# --------------------------------------------------------------------------
@dataclass(frozen=True)
class Literal:
pred: str
args: tuple
negated: bool = False
def __str__(self) -> str:
head = "not " if self.negated else ""
return f"{head}{self.pred}({', '.join(self.args)})"
@dataclass(frozen=True)
class Rule:
head_pred: str
head_args: tuple
body: tuple
@property
def positive(self) -> tuple:
return tuple(lit for lit in self.body if not lit.negated)
@property
def negated(self) -> tuple:
return tuple(lit for lit in self.body if lit.negated)
def __str__(self) -> str:
head = f"{self.head_pred}({', '.join(self.head_args)})"
body = ", ".join(str(lit) for lit in self.body)
return f"{head} :- {body}."
@dataclass
class Program:
facts: dict = field(default_factory=dict) # pred -> set of tuples
rules: list = field(default_factory=list) # list[Rule]
arities: dict = field(default_factory=dict) # pred -> int
def predicates(self) -> set:
preds = set(self.facts)
preds.update(self.arities)
for rule in self.rules:
preds.add(rule.head_pred)
for lit in rule.body:
preds.add(lit.pred)
return preds
def add_fact(self, pred: str, args: tuple) -> None:
self.facts.setdefault(pred, set()).add(tuple(args))
def add_rule(self, rule: Rule) -> None:
self.rules.append(rule)
def text(self) -> str:
lines = []
for pred in sorted(self.facts):
for tup in sorted(self.facts[pred]):
lines.append(f"{pred}({', '.join(tup)}).")
lines.extend(str(r) for r in self.rules)
return "\n".join(lines) + "\n"
# --------------------------------------------------------------------------
# Tokenizer + parser
# --------------------------------------------------------------------------
_TOKEN_RE = re.compile(
r"(?P<ws>[ \t\r\n]+)"
r"|(?P<comment>%[^\n]*)"
r"|(?P<string>\"(?:[^\"\\]|\\.)*\")"
r"|(?P<int>-?\d+)"
r"|(?P<name>[A-Za-z_][A-Za-z0-9_]*)"
r"|(?P<punct>:-|[!(),.])"
)
def tokenize(text: str) -> list:
tokens = []
pos = 0
n = len(text)
while pos < n:
m = _TOKEN_RE.match(text, pos)
if m is None:
line = text.count("\n", 0, pos) + 1
raise ParseError(
f"unexpected character {text[pos]!r} at line {line}, column "
f"{pos - text.rfind(chr(10), 0, pos)}"
)
kind = m.lastgroup
if kind not in ("ws", "comment"):
tokens.append((kind, m.group()))
pos = m.end()
return tokens
class _Parser:
def __init__(self, tokens):
self.toks = tokens
self.i = 0
def peek(self, k: int = 0):
j = self.i + k
return self.toks[j] if j < len(self.toks) else (None, None)
def next(self):
tok = self.peek()
self.i += 1
return tok
def expect(self, kind, value=None):
k, v = self.next()
if k != kind or (value is not None and v != value):
want = value if value is not None else kind
raise ParseError(f"expected {want!r} but found {v!r}")
return v
# ---- grammar ------------------------------------------------------
def parse_term(self):
kind, value = self.next()
if kind in ("string", "int"):
return value
if kind == "name":
if value in ("not", "!"):
raise ParseError(f"{value!r} cannot be used as a term")
if value[0].isupper() or value[0] == "_":
return "?" + value
return value
raise ParseError(f"expected a term but found {value!r}")
def parse_literal(self):
negated = False
kind, value = self.peek()
if (kind == "name" and value == "not") or (kind == "punct" and value == "!"):
self.next()
negated = True
pred = self.expect("name")
if pred in ("not", "!"):
raise ParseError(f"{pred!r} cannot be used as a predicate name")
self.expect("punct", "(")
args = []
if self.peek() == ("punct", ")"):
self.next()
else:
while True:
args.append(self.parse_term())
kind, value = self.peek()
if (kind, value) == ("punct", ","):
self.next()
continue
break
self.expect("punct", ")")
return Literal(pred, tuple(args), negated)
def parse_program(self) -> Program:
program = Program()
while self.peek()[0] is not None:
head = self.parse_literal()
if head.negated:
raise ParseError("a clause head cannot be negated")
if self.peek() == ("punct", ":-"):
self.next()
body = [self.parse_literal()]
while self.peek() == ("punct", ","):
self.next()
body.append(self.parse_literal())
self.expect("punct", ".")
program.add_rule(Rule(head.pred, head.args, tuple(body)))
else:
self.expect("punct", ".")
for term in head.args:
if is_var(term):
raise ParseError(
f"fact {head.pred}/{len(head.args)} contains the "
f"variable {term}; facts must be ground"
)
program.add_fact(head.pred, head.args)
return program
def parse_program(text: str) -> Program:
"""Parse Datalog source text into a Program."""
return _Parser(tokenize(text)).parse_program()
def load(path: str) -> Program:
with open(path, "r", encoding="utf-8") as fh:
return parse_program(fh.read())
# --------------------------------------------------------------------------
# Static checks
# --------------------------------------------------------------------------
def check_arity(program: Program) -> None:
"""Every predicate must be used with one arity throughout."""
seen = dict(program.arities)
def note(pred, n, where):
if pred in seen and seen[pred] != n:
raise ArityError(
f"predicate {pred!r} used with arity {n} in {where} "
f"but arity {seen[pred]} elsewhere"
)
seen[pred] = n
for pred, tuples in program.facts.items():
for tup in tuples:
note(pred, len(tup), f"a fact of {pred!r}")
for rule in program.rules:
note(rule.head_pred, len(rule.head_args), f"the head of {rule}")
for lit in rule.body:
note(lit.pred, len(lit.args), f"the body of {rule}")
def check_safety(program: Program) -> None:
"""Range restriction: head and negated-literal variables must be bound by a
positive body literal, and every rule needs at least one positive literal."""
for rule in program.rules:
if not rule.positive:
raise SafetyError(f"rule {rule} has no positive body literal (unsafe)")
bound = set()
for lit in rule.positive:
for term in lit.args:
if is_var(term) and not is_anon(term):
bound.add(term)
for term in rule.head_args:
if is_anon(term):
raise SafetyError(f"anonymous variable in the head of {rule}")
if is_var(term) and term not in bound:
raise SafetyError(
f"head variable {term} of {rule} is not bound by a positive "
f"body literal (unsafe)"
)
for lit in rule.negated:
for term in lit.args:
if is_var(term) and not is_anon(term) and term not in bound:
raise SafetyError(
f"variable {term} in the negated literal {lit} of {rule} "
f"is not bound by a positive body literal (unsafe)"
)
def _tarjan_scc(nodes, adj) -> list:
"""Iterative Tarjan SCC. Returns components as sorted lists, deterministic."""
index: dict = {}
low: dict = {}
on_stack: dict = {}
stack: list = []
out: list = []
counter = [0]
for root in nodes:
if root in index:
continue
index[root] = low[root] = counter[0]
counter[0] += 1
stack.append(root)
on_stack[root] = True
work = [(root, iter(sorted(adj.get(root, ()))))]
while work:
node, it = work[-1]
advanced = False
for nxt in it:
if nxt not in index:
index[nxt] = low[nxt] = counter[0]
counter[0] += 1
stack.append(nxt)
on_stack[nxt] = True
work.append((nxt, iter(sorted(adj.get(nxt, ())))))
advanced = True
break
if on_stack.get(nxt):
low[node] = min(low[node], index[nxt])
if advanced:
continue
work.pop()
if work:
parent = work[-1][0]
low[parent] = min(low[parent], low[node])
if low[node] == index[node]:
comp = []
while True:
top = stack.pop()
on_stack[top] = False
comp.append(top)
if top == node:
break
out.append(sorted(comp))
return out
def stratify(program: Program) -> list:
"""Split predicates into strata (a list of frozensets, dependencies first).
Raises StratificationError when a negated literal sits inside a recursive
cycle, i.e. when head and negated body predicate share an SCC.
"""
preds = sorted(program.predicates())
adj = {p: set() for p in preds}
neg_edges = []
for rule in program.rules:
for lit in rule.body:
adj[rule.head_pred].add(lit.pred)
if lit.negated:
neg_edges.append((rule.head_pred, lit.pred))
comps = _tarjan_scc(preds, adj)
comp_of = {}
for idx, comp in enumerate(comps):
for pred in comp:
comp_of[pred] = idx
for head, body in neg_edges:
if comp_of[head] == comp_of[body]:
members = ", ".join(comps[comp_of[head]])
raise StratificationError(
f"negation is not stratified: {head} negates {body} but they are "
f"mutually recursive (SCC: {members})"
)
# condensation: edge comp(p) -> comp(q) means "p depends on q", so q must be
# computed first. Kahn over the reversed relation.
must_precede = {i: set() for i in range(len(comps))}
for pred in preds:
for dep in adj[pred]:
a, b = comp_of[pred], comp_of[dep]
if a != b:
must_precede[a].add(b)
indeg = {i: len(must_precede[i]) for i in must_precede}
dependents = {i: set() for i in must_precede}
for i, deps in must_precede.items():
for dep in deps:
dependents[dep].add(i)
ready = sorted((i for i, d in indeg.items() if d == 0), key=lambda i: comps[i][0])
order = []
while ready:
cur = ready.pop(0)
order.append(cur)
for nxt in sorted(dependents[cur], key=lambda i: comps[i][0]):
indeg[nxt] -= 1
if indeg[nxt] == 0:
ready.append(nxt)
ready.sort(key=lambda i: comps[i][0])
if len(order) != len(comps): # pragma: no cover - condensation is a DAG
raise StratificationError("internal error: condensation is not a DAG")
return [frozenset(comps[i]) for i in order]
# --------------------------------------------------------------------------
# Evaluation
# --------------------------------------------------------------------------
def _match(pattern, tup, env):
"""Unify a literal's argument pattern against a tuple. Returns a new env or
None. Anonymous variables never bind. Bindings made while walking this
literal are visible to its later arguments, so `p(X, X)` really does mean
"both arguments equal"."""
cur = env
copied = False
for term, value in zip(pattern, tup):
if is_var(term):
if is_anon(term):
continue
name = var_name(term)
if name in cur:
if cur[name] != value:
return None
else:
if not copied:
cur = dict(env)
copied = True
cur[name] = value
elif term != value:
return None
return cur
def _subst(term, env):
if is_var(term):
if is_anon(term):
raise SafetyError("anonymous variable cannot be substituted")
return env[var_name(term)]
return term
def _negation_fails(lit, env, rels):
"""True when the negated literal matches some tuple (so the body fails)."""
rel = rels.get(lit.pred, ())
if not rel:
return False
for tup in rel:
if _match(lit.args, tup, env) is not None:
return True
return False
def _order_pos(rule, R, rels):
"""Positive body literals, smallest relation first (deterministic)."""
def size(lit):
if lit.pred in R:
return len(R[lit.pred])
return len(rels.get(lit.pred, ()))
return sorted(rule.positive, key=lambda lit: (size(lit), lit.pred))
def _join(rule, pos, R, rels, override, stats):
"""Enumerate satisfying environments for the positive body literals."""
def relation_for(idx, lit):
if override and idx in override:
return override[idx]
if lit.pred in R:
return R[lit.pred]
return rels.get(lit.pred, ())
def rec(idx, env):
if idx == len(pos):
for lit in rule.negated:
if _negation_fails(lit, env, rels):
return
yield env
return
lit = pos[idx]
for tup in sorted(relation_for(idx, lit)):
stats["bindings"] += 1
env2 = _match(lit.args, tup, env)
if env2 is not None:
yield from rec(idx + 1, env2)
yield from rec(0, {})
def _derive_one(rule, pos, R, rels, override, out, stats):
head_pred = rule.head_pred
for env in _join(rule, pos, R, rels, override, stats):
stats["rule_runs"] += 1
tup = tuple(_subst(term, env) for term in rule.head_args)
if tup in R[head_pred] or tup in out[head_pred]:
continue
out[head_pred].add(tup)
stats["derived"] += 1
def _derive(rules_here, R, rels, wave, override, stats):
out = {pred: set() for pred in wave}
for rule in rules_here:
pos = _order_pos(rule, R, rels)
_derive_one(rule, pos, R, rels, override, out, stats)
return out
def _eval_stratum(program, wave, rels, stats):
rules_here = [r for r in program.rules if r.head_pred in wave]
if not rules_here:
return
R = {pred: rels[pred] for pred in wave} # live sets, mutated in place
# Round 0 is naive: it also fires rules whose positive body lives entirely in
# lower strata (such a rule can never see a same-stratum delta).
new = _derive(rules_here, R, rels, wave, None, stats)
for pred, added in new.items():
R[pred] |= added
delta = {pred: set(added) for pred, added in new.items()}
while any(delta.values()):
stats["rounds"] += 1
new = {pred: set() for pred in wave}
for rule in rules_here:
pos = _order_pos(rule, R, rels)
for idx, lit in enumerate(pos):
if lit.pred in wave and delta.get(lit.pred):
_derive_one(
rule, pos, R, rels, {idx: delta[lit.pred]}, new, stats
)
for pred, added in new.items():
R[pred] |= added
delta = new
def evaluate(program: Program, stats: dict = None, check: bool = True) -> dict:
"""Semi-naive bottom-up evaluation. Returns {pred: frozenset(tuples)}."""
if check:
check_arity(program)
check_safety(program)
strata = stratify(program)
if stats is None:
stats = {}
stats.update(rounds=0, bindings=0, derived=0, rule_runs=0)
rels = {pred: set(program.facts.get(pred, ())) for pred in program.predicates()}
for wave in strata:
_eval_stratum(program, wave, rels, stats)
return {pred: frozenset(tuples) for pred, tuples in rels.items()}
# --------------------------------------------------------------------------
# Output
# --------------------------------------------------------------------------
def format_relations(rels: dict, preds=None) -> str:
lines = []
for pred in sorted(rels) if preds is None else list(preds):
tuples = sorted(rels.get(pred, ()))
lines.append(f"{pred}: {len(tuples)} tuple(s)")
for tup in tuples:
lines.append(f" {pred}({', '.join(tup)}).")
return "\n".join(lines)
def main(argv=None) -> int:
ap = argparse.ArgumentParser(
prog="datalog.py",
description="Evaluate a Datalog program (stratified negation, semi-naive).",
)
ap.add_argument("program", help="path to a Datalog source file")
ap.add_argument("-q", "--query", action="append", default=None,
help="only report this predicate (repeatable)")
ap.add_argument("--text", action="store_true", help="human-readable output")
ap.add_argument("--stats", action="store_true", help="include run statistics")
args = ap.parse_args(argv)
try:
program = load(args.program)
stats = {}
rels = evaluate(program, stats)
strata = stratify(program)
except DatalogError as exc:
print(f"error: {type(exc).__name__}: {exc}", file=sys.stderr)
return 2
except OSError as exc:
print(f"error: {exc}", file=sys.stderr)
return 2
if args.text:
print(format_relations(rels, args.query))
if args.stats:
print(json.dumps(stats, sort_keys=True))
return 0
payload = {
"predicates": {
pred: [list(t) for t in sorted(rels[pred])]
for pred in (sorted(rels) if args.query is None else args.query)
},
"strata": [sorted(wave) for wave in strata],
}
if args.stats:
payload["stats"] = stats
print(json.dumps(payload, sort_keys=True, indent=2))
return 0
if __name__ == "__main__": # pragma: no cover
sys.exit(main())