| """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() |
| |
| 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 |
| |
| 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.) |
| |
| 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) |
|
|
