"""Checks for the new architectural connection, action contract, and checkpoints.""" from dataclasses import replace import json import pytest import torch from lightnest import rv32 from lightnest.relation_patch import ( RelationConfig, RelationPatchNetwork, RelationPredictor, definition_use_graph, encode_batch, load_records, load_relation, supervised_loss, ) def assemble(words): return b''.join(x.to_bytes(4, 'little') for x in words) def far_pair(): words = [rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2')] words += [rv32.addi('x0', 'x0', 31 + i) for i in range(8)] q = len(words) * 4 + 3 words += [rv32.add('x4', 'x1', 'x2')] words += [rv32.addi('x0', 'x0', 211 + i) for i in range(8)] pair = [assemble(words + [rv32.addi('x10', reg, 0), rv32.EBREAK]) for reg in ('x3', 'x4')] return pair, (11, q) def test_definition_edges_obey_overwrites_zero_and_implicit_return(): code = assemble([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2'), rv32.sub('x3', 'x2', 'x1'), rv32.addi('x0', 'x3', 0), rv32.addi('x10', 'x3', 0), rv32.EBREAK]) graph = definition_use_graph(code) assert graph.instruction_count == 7 assert set(graph.edges) == {(2, 0), (2, 1), (3, 0), (3, 1), (4, 3), (5, 3), (6, 5)} assert (5, 2) not in graph.edges @pytest.mark.parametrize('code', [b'', b'\x13', assemble([rv32.beq('x1', 'x2', 4)]), assemble([rv32.EBREAK, rv32.addi('x0', 'x0', 0)])]) def test_unsupported_structure_is_explicit(code): with pytest.raises(ValueError): definition_use_graph(code) def test_local_pointer_is_invariant_but_graph_pointer_reads_remote_return(): torch.set_num_threads(2) torch.manual_seed(137) pair, (p, q) = far_pair() # Exactly identical 17-byte neighborhoods at both candidate positions. for pos in (p, q): assert pair[0][pos - 8:pos + 9] == pair[1][pos - 8:pos + 9] graph = RelationPatchNetwork(RelationConfig(mode='dataflow')).eval() local = RelationPatchNetwork(replace(graph.config, mode='local')).eval() local.load_state_dict(graph.state_dict()) assert sum(p.numel() for p in graph.parameters()) == sum(p.numel() for p in local.parameters()) encoded = encode_batch(['change the returned add to sub'] * 2, pair, graph.config) with torch.inference_mode(): ls, _ = local(*encoded) gs, _ = graph(*encoded) local_change = (ls[0, p] - ls[0, q]) - (ls[1, p] - ls[1, q]) graph_change = (gs[0, p] - gs[0, q]) - (gs[1, p] - gs[1, q]) assert float(local_change.abs()) < 1e-7 # This proves a connection exists, not that random weights solve the task. assert float(graph_change.abs()) > 1e-6 def test_dynamic_batch_padding_preserves_inference_and_deeper_passes(): torch.manual_seed(17) model = RelationPatchNetwork(RelationConfig(width=32)).eval() pair, _ = far_pair() short = assemble([rv32.addi('x10', 'x0', 9), rv32.EBREAK]) with torch.inference_mode(): first = model(*encode_batch(['no changes'], [short], model.config)) together = model(*encode_batch(['no changes', 'replace the returned add with sub'], [short, pair[0]], model.config)) deeper = model(*encode_batch(['no changes'], [short], model.config), message_steps=8) for x, y in zip(first, together): torch.testing.assert_close(x[0], y[0], atol=2e-6, rtol=2e-5) assert deeper[0].shape == (1, 1025) assert not torch.equal(first[0], deeper[0]) def test_set_target_loss_accepts_both_orders_and_explicit_stop(): scores = torch.zeros(2, 9, requires_grad=True) bits = torch.zeros(2, 8, 8, requires_grad=True) rows = [{'edits': [{'offset': 1, 'value': 65}, {'offset': 5, 'value': 2}]}, {'edits': []}] loss, pointer_loss, byte_loss = supervised_loss(scores, bits, rows, none_index=8) expected = (-torch.log(torch.tensor(2 / 9)) - torch.log(torch.tensor(1 / 9))) / 2 torch.testing.assert_close(pointer_loss, expected) torch.testing.assert_close(byte_loss, torch.tensor(2.).log()) loss.backward() assert scores.grad[0, 1] < 0 and scores.grad[0, 5] < 0 and scores.grad[1, 8] < 0 assert torch.count_nonzero(bits.grad[0, 1]) == 8 assert torch.count_nonzero(bits.grad[0, 5]) == 8 assert torch.count_nonzero(bits.grad[1]) == 0 def test_checkpoint_roundtrip_and_reject_wrong_architecture_or_nonfinite(tmp_path): model = RelationPatchNetwork(RelationConfig(width=32, mode='local')) path = tmp_path / 'model.pt' model.save(path) restored = load_relation(path) assert restored.config == model.config for a, b in zip(model.parameters(), restored.parameters()): torch.testing.assert_close(a, b, rtol=0., atol=0.) saved = torch.load(path, weights_only=True) saved['architecture'] = 'absolute-byte-control' torch.save(saved, path) with pytest.raises(ValueError, match='architecture'): load_relation(path) saved['architecture'] = model.architecture saved['state_dict']['pointer.weight'][0, 0] = float('nan') torch.save(saved, path) with pytest.raises(ValueError, match='tensor'): load_relation(path) def test_adapter_uses_actual_byte_precondition_and_explicit_stop(): from lightnest.iterative_patch import run_iterative_patch class FixedNetwork: config = RelationConfig() none_index = 1024 def eval(self): return self def __call__(self, command, binary, relations, message_steps=None): scores = torch.full((1, 1025), -40.) bits = torch.full((1, 1024, 8), -40.) # A fixture action tests the adapter only; this is not a trained model. value = int(binary[0, 3]) - 1 if value == 64: scores[0, 1024] = 40. else: scores[0, 3] = 40. bits[0, 3, 6] = 40. return scores, bits code = assemble([rv32.add('x3', 'x1', 'x2'), rv32.addi('x10', 'x3', 0), rv32.EBREAK]) predictor = RelationPredictor(FixedNetwork()) decision = predictor('change add to sub', code) assert decision.edits[0].offset == 3 assert decision.edits[0].expected == code[3:4] assert decision.edits[0].replacement == bytes([64]) result = run_iterative_patch('change add to sub', code, predictor, max_steps=2) assert result.final[3] == 64 assert result.final[:3] == code[:3] and result.final[4:] == code[4:] def test_dataset_checks_do_not_silently_miss_an_edit(tmp_path): binary = assemble([rv32.addi('x10', 'x0', 9), rv32.EBREAK]) path = tmp_path / 'rows.jsonl' row = {'id': 'case', 'command': 'do not change', 'binary': binary.hex(), 'target': binary.hex(), 'edits': []} path.write_text(json.dumps(row) + '\n') assert load_records(path)[0]['edits'] == [] row['target'] = bytes([binary[0] ^ 1]).hex() + binary[1:].hex() path.write_text(json.dumps(row) + '\n') with pytest.raises(ValueError, match='disagree'): load_records(path)