check_indexed_plans.py

#!/usr/bin/env python3
"""Finite affine plan checking, conditional on named primitive contracts.

This is a handwritten contract checker, not OCaml extraction or Rocq proof.
Run with --source PATH to pin the inspected upstream files as well.
"""
from dataclasses import dataclass
from pathlib import Path
import argparse
import hashlib
import json
from itertools import product

PARAMETERS = ('n', 'p', 'r', 's')


class Rejected(ValueError):
    pass


def require(ok, message):
    if not ok:
        raise Rejected(message)


@dataclass(frozen=True)
class Affine:
    coefficients: tuple

    @staticmethod
    def read(x):
        if isinstance(x, int):
            return Affine((x, 0, 0, 0, 0))
        require(isinstance(x, dict) and set(x) <= {'const', *PARAMETERS},
                'unknown affine parameter')
        require(all(type(v) is int for v in x.values()), 'noninteger coefficient')
        return Affine(tuple(x.get(k, 0) for k in ('const', *PARAMETERS)))

    def __add__(self, other):
        if isinstance(other, int):
            other = Affine.read(other)
        return Affine(tuple(a + b for a, b in zip(self.coefficients, other.coefficients)))

    def __sub__(self, other):
        if isinstance(other, int):
            other = Affine.read(other)
        return Affine(tuple(a - b for a, b in zip(self.coefficients, other.coefficients)))

    def nonnegative(self):
        return all(c >= 0 for c in self.coefficients)

    def at(self, values):
        return self.coefficients[0] + sum(c * v for c, v in zip(self.coefficients[1:], values))


@dataclass(frozen=True)
class Type:
    sort: str
    level: Affine
    size: Affine | None = None
    arity: int | None = None
    kind: str | None = None
    left: int | None = None
    right: int | None = None


def read_type(t):
    require(t['sort'] in ('Stored', 'Buf', 'Chain', 'Semi'), 'unknown sort')
    level = Affine.read(t['level'])
    require(level.nonnegative(), 'level is not uniformly natural')
    if t['sort'] == 'Buf':
        size = Affine.read(t['size'])
        require(size.nonnegative(), 'occupancy is not uniformly natural')
        return Type('Buf', level, size)
    if t['sort'] == 'Chain':
        require(t['arity'] in (0, 1, 2), 'invalid arity')
        require(t['kind'] in ('only', 'left', 'right'), 'invalid kind')
        require(type(t['left']) is int and 0 <= t['left'] < 16, 'invalid left color')
        require(type(t['right']) is int and 0 <= t['right'] < 16, 'invalid right color')
        return Type('Chain', level, arity=t['arity'], kind=t['kind'],
                    left=t['left'], right=t['right'])
    return Type(t['sort'], level)


@dataclass(frozen=True)
class Value:
    type: Type
    word: tuple


class Checker:
    def __init__(self, plan):
        self.values = {}
        self.substitution = {}
        self.module_calls = 0
        self.fresh_records = 0
        self.depth = 0
        for name, root in plan['inputs'].items():
            self.values[name] = Value(read_type(root['type']), tuple(root['word']))
            require(type(root.get('path_depth', 0)) is int and root.get('path_depth', 0) >= 0,
                    'invalid declared access path')
            self.depth = max(self.depth, root.get('path_depth', 0))

    def take(self, name, sort=None):
        require(name in self.values, 'unavailable or forward reference')
        v = self.values[name]
        if sort:
            require(v.type.sort == sort, f'expected {sort}, got {v.type.sort}')
        return v

    def put(self, name, t, word):
        require(name not in self.values, 'duplicate binding')
        self.values[name] = Value(t, tuple(word))

    def normalize(self, word):
        answer = ()
        for atom in word:
            answer += self.normalize(self.substitution[atom]) if atom in self.substitution else (atom,)
        return answer

    def instruction(self, i):
        op = i['op']
        if op == 'empty':
            t = read_type({'sort': 'Chain', 'level': i['level'], 'arity': 0,
                           'kind': i['kind'], 'left': i['left'], 'right': i['right']})
            self.put(i['out'], t, ())
            self.fresh_records += 1
        elif op == 'ground':
            require(i['letter'] in ('a', 'b'), 'unknown payload letter')
            self.put(i['out'], Type('Stored', Affine.read(0)), (i['letter'],))
            self.fresh_records += 1
        elif op == 'small':
            b = self.take(i['buffer'], 'Buf')
            require((b.type.size - 3).nonnegative(), 'Small needs at least 3 entries')
            self.put(i['out'], Type('Stored', b.type.level + 1), b.word)
            self.fresh_records += 1
        elif op == 'big':
            p, c, s = self.take(i['prefix'], 'Buf'), self.take(i['child'], 'Chain'), self.take(i['suffix'], 'Buf')
            require(p.type.level == s.type.level, 'Big buffer level mismatch')
            require(c.type.level == p.type.level + 1, 'Big child level mismatch')
            require(c.type.kind == 'only', 'Big requires only-kind child')
            require((p.type.size - 3).nonnegative() and (s.type.size - 3).nonnegative(),
                    'Big needs buffers of at least 3 entries')
            self.put(i['out'], Type('Stored', p.type.level + 1), p.word + c.word + s.word)
            self.fresh_records += 1
        elif op == 'pack':
            c = self.take(i['chain'], 'Chain')
            require(c.type.kind == 'only', 'Semi requires only-kind chain')
            self.put(i['out'], Type('Semi', c.type.level), c.word)
            self.fresh_records += 1
        elif op in ('push', 'concat'):
            if op == 'push':
                x, c = self.take(i['stored'], 'Stored'), self.take(i['child'], 'Semi')
            else:
                x, c = self.take(i['left'], 'Semi'), self.take(i['right'], 'Semi')
            require(x.type.level == c.type.level, 'chain operation level mismatch')
            self.put(i['out'], Type('Semi', c.type.level), x.word + c.word)
            self.module_calls += 1
        elif op in ('inject3', 'inject_vector'):
            p, v = self.take(i['prefix'], 'Buf'), self.take(i['items'], 'Buf')
            require(p.type.level == v.type.level, 'injection level mismatch')
            if op == 'inject3':
                require(v.type.size == Affine.read(3), 'inject3 requires exactly 3 entries')
            else:
                require(v.type.size in tuple(Affine.read(k) for k in range(3)),
                        'inject_vector requires fixed 0, 1, or 2 entries')
            self.put(i['out'], Type('Buf', p.type.level, p.type.size + v.type.size), p.word + v.word)
            self.module_calls += 1
        elif op == 'split3':
            b = self.take(i['buffer'], 'Buf')
            require(len(b.word) == 1 and b.word[0] not in self.substitution,
                    'split requires fresh opaque buffer word')
            if i['case'] == 'short':
                require(i['remainder'] in (0, 1, 2), 'invalid short case')
                size = Affine.read(i['remainder'])
            else:
                require(i['case'] == 'long', 'unknown split case')
                size = Affine.read(i['remainder'])
                require((size - 3).nonnegative(), 'long remainder needs at least 3 entries')
            require(b.type.size == size + 3, 'split occupancy equation fails')
            uword = (i['three'],)
            rword = () if size == Affine.read(0) else (i['rest'],)
            used_atoms = {'a', 'b'} | {a for v in self.values.values() for a in v.word}
            used_atoms |= {a for w in self.substitution.values() for a in w}
            require(b.word[0] not in ('a', 'b') and i['three'] != i['rest']
                    and i['three'] not in used_atoms and i['rest'] not in used_atoms,
                    'split word-name collision')
            self.substitution[b.word[0]] = uword + rword
            self.put(i['three'], Type('Buf', b.type.level, Affine.read(3)), uword)
            self.put(i['rest'], Type('Buf', b.type.level, size), rword)
            self.module_calls += 1
        else:
            raise Rejected('unknown operation')

    def check(self, plan):
        for instruction in plan['instructions']:
            self.instruction(instruction)
        outputs = []
        for spec in plan['outputs']:
            v = self.take(spec['value'])
            if 'type' in spec:
                require(v.type == read_type(spec['type']), 'output type mismatch')
            if 'min_size' in spec:
                require(v.type.sort == 'Buf' and (v.type.size - spec['min_size']).nonnegative(),
                        'output occupancy refinement fails')
            outputs.append(v.word)
        if 'target_word' in plan:
            require(self.normalize(sum(outputs, ())) == self.normalize(tuple(plan['target_word'])),
                    'word identity fails')
        return {'module_calls': self.module_calls, 'fresh_records': self.fresh_records,
                'declared_old_path_depth': self.depth}


def certify(plan):
    return Checker(plan).check(plan)


def buffer(size, word, depth=0, level=None):
    return {'type': {'sort': 'Buf', 'level': level or {'n': 1}, 'size': size},
            'word': [word], 'path_depth': depth}


def repair_plan(kind, short):
    rest = short if short is not None else {'const': 3, 'r': 1}
    head_size = short + 3 if short is not None else {'const': 6, 'r': 1}
    roots = {'P': buffer({'const': 5, 'p': 1}, 'P'), 'H': buffer(head_size, 'H', 1),
             'D': {'type': {'sort': 'Semi', 'level': {'const': 1, 'n': 1}}, 'word': ['D']}}
    code = []
    target = ['P', 'H', 'D']
    child = 'D'
    if kind == 'Big':
        roots['S'] = buffer({'const': 3, 's': 1}, 'S', 1)
        roots['T'] = {'type': {'sort': 'Chain', 'level': {'const': 1, 'n': 1},
                              'arity': 2, 'kind': 'only', 'left': 1, 'right': 8},
                      'word': ['T'], 'path_depth': 1}
        code += [{'op': 'small', 'buffer': 'S', 'out': 's1'},
                 {'op': 'push', 'stored': 's1', 'child': 'D', 'out': 'd1'},
                 {'op': 'pack', 'chain': 'T', 'out': 't1'},
                 {'op': 'concat', 'left': 't1', 'right': 'd1', 'out': 'd2'}]
        child = 'd2'
        target = ['P', 'H', 'T', 'S', 'D']
    code += [{'op': 'split3', 'buffer': 'H', 'case': 'short' if short is not None else 'long',
              'remainder': rest, 'three': 'U', 'rest': 'V'},
             {'op': 'inject3', 'prefix': 'P', 'items': 'U', 'out': 'p1'}]
    prefix = 'p1'
    if short is not None:
        code += [{'op': 'inject_vector', 'prefix': 'p1', 'items': 'V', 'out': 'p2'}]
        prefix = 'p2'
    else:
        code += [{'op': 'small', 'buffer': 'V', 'out': 'v1'},
                 {'op': 'push', 'stored': 'v1', 'child': child, 'out': 'd3'}]
        child = 'd3'
    return {'inputs': roots, 'instructions': code,
            'outputs': [{'value': prefix, 'min_size': 8},
                        {'value': child, 'type': {'sort': 'Semi', 'level': {'const': 1, 'n': 1}}}],
            'target_word': target}


SOURCE_HASHES = {
    'theory/Cadeque/Types.v': '4beb995bd9634c6cc275aa053e009065a7ff699e2ad7b04fd2821df97bcd9e23',
    'theory/Cadeque/Models.v': 'd19612a21f4e506aedb31657393c437a07d1f2f8e5ed4ebcea4e408d6b05402f',
    'lib/cadeque_core.ml': '2a7e36a50e9cec5f3cc06226922ab26e039e1fa0dc1188e879e594f320ae003e',
}


def run():
    import copy
    accepted = {f'{kind}/{k if k is not None else "long"}': certify(repair_plan(kind, k))
                for kind in ('Small', 'Big') for k in (0, 1, 2, None)}
    failures = {}
    def rejects(name, plan):
        try:
            certify(plan)
        except Rejected as e:
            failures[name] = str(e)
        else:
            raise AssertionError(f'negative control accepted: {name}')
    bad = repair_plan('Small', 2)
    bad['instructions'][-1:] = [{'op': 'small', 'buffer': 'V', 'out': 'badsmall'},
                               {'op': 'push', 'stored': 'badsmall', 'child': 'D', 'out': 'badchild'}]
    bad['outputs'][0]['value'] = 'p1'
    bad['outputs'][1]['value'] = 'badchild'
    rejects('two-entry-Small-word-correct', bad)
    bad = repair_plan('Small', None)
    bad['instructions'][-2:] = [{'op': 'push', 'stored': 'V', 'child': 'D', 'out': 'd3'}]
    rejects('omitted-Small-word-correct', bad)
    bad = repair_plan('Big', None)
    bad['inputs']['S']['type']['level'] = {'const': 1, 'n': 1}
    rejects('wrong-suffix-level', bad)
    bad = repair_plan('Big', None)
    bad['inputs']['T']['type']['kind'] = 'left'
    rejects('wrong-child-kind', bad)
    bad = repair_plan('Big', None)
    bad['inputs']['T']['type']['left'] = 16
    rejects('invalid-color', bad)
    bad = repair_plan('Big', None)
    bad['inputs']['T']['type']['arity'] = 3
    rejects('invalid-arity', bad)
    bad = repair_plan('Small', 1)
    bad['inputs']['P']['type']['size'] = {'const': 5, 'p': -1}
    rejects('eventually-negative-occupancy', bad)
    bad = repair_plan('Small', None)
    bad['instructions'][-1]['child'] = 'd3'
    rejects('forward-reference', bad)
    bad = repair_plan('Big', None)
    bad['instructions'][3]['left'], bad['instructions'][3]['right'] = 'd1', 't1'
    rejects('reversed-concat', bad)
    # Empty is specialized at each explicit use, independently of stored rank.
    empty_samples = 0
    for level in (0, {'n': 1}, {'const': 100, 'n': 1}):
        for kind, lc, rc in product(('only', 'left', 'right'), range(16), range(16)):
            certify({'inputs': {}, 'instructions': [{'op': 'empty', 'out': 'e',
                      'level': level, 'kind': kind, 'left': lc, 'right': rc}],
                     'outputs': [{'value': 'e'}], 'target_word': []})
            empty_samples += 1
    poly = {'inputs': {'P': buffer({'const': 3, 'p': 1}, 'P'),
                       'S': buffer({'const': 3, 's': 1}, 'S')},
            'instructions': [{'op': 'empty', 'out': 'e', 'level': {'const': 1, 'n': 1},
                              'kind': 'only', 'left': 0, 'right': 15},
                             {'op': 'big', 'prefix': 'P', 'child': 'e', 'suffix': 'S', 'out': 'b'}],
            'outputs': [{'value': 'b', 'type': {'sort': 'Stored', 'level': {'const': 1, 'n': 1}}}],
            'target_word': ['P', 'S']}
    certify(poly)
    bad = copy.deepcopy(poly)
    bad['instructions'][0]['kind'] = 'left'
    rejects('Big-empty-wrong-kind', bad)
    bad = copy.deepcopy(poly)
    bad['instructions'][0]['level'] = {'n': 1}
    rejects('Big-empty-wrong-level', bad)
    # Finite differential corroboration for affine obligation rejection.
    arithmetic_cases = 0
    for coeffs in product(range(-2, 3), repeat=3):
        a = Affine((coeffs[0], coeffs[1], coeffs[2], 0, 0))
        brute = all(a.at((n, p, 0, 0)) >= 0 for n, p in product(range(9), repeat=2))
        assert a.nonnegative() == brute
        arithmetic_cases += 1
    return {'scope': 'Generic annotated straight-line contract-plan checker with independent natural affine parameters. Eight finite source-clause subcases checked symbolically; finite Empty/metadata/arithmetic samples and mutation rejection. No OCaml execution/refinement, Rocq build, proof-assistant verification, general type inference, unrestricted plan synthesis, primitive implementation certificate, or iterative closure proof.',
            'accepted_repair_plans': accepted, 'rejected_controls': failures,
            'Empty_specializations_sampled': empty_samples,
            'finite_arithmetic_differential_cases': arithmetic_cases,
            'polymorphic_Empty_Big_plan': 'accepted',
            'closure_witness': 'P has five entries, z=Small(U) with three entries, green singleton child holding z. Pop leaves Empty; short repair0 leaves Empty, outside the original nonempty-green-child input view.',
            'source_commit': '503c5d3ff7f9a20bad1cd7a8f8d7abf4e635d3a0',
            'source_hashes': SOURCE_HASHES}


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', type=Path)
    parser.add_argument('--plan', type=Path, help='check one supplied JSON plan')
    args = parser.parse_args()
    if args.source:
        for path, expected in SOURCE_HASHES.items():
            require(hashlib.sha256((args.source / path).read_bytes()).hexdigest() == expected,
                    f'source hash mismatch: {path}')
    if args.plan:
        print(json.dumps(certify(json.loads(args.plan.read_text())), indent=2))
    else:
        print(json.dumps(run(), indent=2))