| import random |
| import unittest |
| from unittest.mock import patch |
| import torch |
| from lightnest import rv32 |
| from lightnest.learned_patch import MAX_WINDOW |
| from lightnest.sequential_patch_rl import arithmetic, make_task, passes, rollout, terminal_reward |
|
|
| class FixedPolicy(torch.nn.Module): |
| def __init__(self, position): |
| super().__init__(); self.position=position; self.calls=0 |
| self.weight=torch.nn.Parameter(torch.tensor(1.)) |
| def forward(self,c,b): |
| self.calls+=1 |
| scores=torch.full((len(c),257),-30.)+self.weight*0 |
| scores[:,self.position]=self.weight*30 |
| return scores,torch.ones((len(c),256,8))*self.weight |
|
|
| class SequentialTests(unittest.TestCase): |
| def test_oracle(self): |
| self.assertEqual(arithmetic((4,2,1),'sub'),1) |
| self.assertEqual(arithmetic((2147483647,1,1),'add'),-2147483647) |
| def test_generator_bounds_and_independent_execution(self): |
| for split in ('train','unseen_bias','three_op_combined'): |
| for seed in range(40): |
| t=make_task(random.Random(seed),split) |
| self.assertLessEqual(len(t.binary),MAX_WINDOW) |
| |
| target='sub' if 'to sub' in t.command or 'with sub' in t.command else 'add' |
| words=[int.from_bytes(t.binary[i:i+4],'little') for i in range(0,len(t.binary),4)] |
| words=[((w & ~(127<<25)) | ((32 if target=='sub' else 0)<<25)) if w&127==51 else w for w in words] |
| repaired=b''.join(w.to_bytes(4,'little') for w in words) |
| self.assertTrue(passes(repaired,t.cases)) |
| def test_none_and_budget(self): |
| t=make_task(random.Random(3)); p=FixedPolicy(MAX_WINDOW) |
| out,edits,h,_=rollout(p,[t.command],[t.binary]) |
| self.assertEqual(out,[t.binary]); self.assertEqual(edits,[0]); self.assertEqual(p.calls,1) |
| p=FixedPolicy(0) |
| with patch('lightnest.sequential_patch_rl.rv32.simulate',side_effect=AssertionError('oracle reached policy')): |
| _,_,h,lp=rollout(p,[t.command],[t.binary],sample=True) |
| self.assertEqual(p.calls,3); self.assertEqual(len(h[0]),3) |
| lp.sum().backward(); self.assertIsNotNone(p.weight.grad) |
| def test_input_bounds(self): |
| p=FixedPolicy(0) |
| for n in (0,4): |
| with self.assertRaises(ValueError): rollout(p,['x'],[b'1234'],max_proposals=n) |
| with self.assertRaises(ValueError): rollout(p,['x'],[b'x'*257]) |
| def test_terminal_reward(self): |
| binary=b''.join(w.to_bytes(4,'little') for w in (rv32.addi('x10','x0',7),rv32.EBREAK)) |
| cases=[([1,2],7)] |
| self.assertEqual(terminal_reward(binary,binary,cases,0),1) |
| self.assertEqual(terminal_reward(binary,binary,cases,1),-1) |
| self.assertEqual(terminal_reward(binary,binary,[([1,2],8)],0),-1) |
| with self.assertRaises(ValueError): terminal_reward(binary,binary,[],0) |
|
|
| if __name__=='__main__': unittest.main() |
|
|
