| import unittest |
| import torch |
| from lightnest import rv32 |
| from lightnest.learned_patch import MAX_WINDOW |
| from lightnest.sequential_stop_rl import rollout, terminal_reward |
|
|
| class Sequence(torch.nn.Module): |
| def __init__(self, positions): |
| super().__init__(); self.positions=positions; self.calls=0 |
| def forward(self,c,b): |
| scores=torch.full((len(c),257),-100.) |
| scores[:,self.positions[self.calls]]=100.; self.calls+=1 |
| return scores,torch.ones((len(c),256,8))*100 |
|
|
| class StopTests(unittest.TestCase): |
| def test_three_edits_then_none(self): |
| p=Sequence([0,1,2,MAX_WINDOW]); out,e,h,lp=rollout(p,['test'],[b'1234']) |
| self.assertEqual(out,[bytes([255,255,255,52])]);self.assertEqual(e,[3]);self.assertEqual(len(h[0]),4) |
| def test_fourth_edit_is_not_applied(self): |
| p=Sequence([0,1,2,3]);out,e,h,_=rollout(p,['test'],[b'1234']) |
| self.assertEqual(out,[bytes([255,255,255,52])]);self.assertEqual(e,[3]);self.assertEqual(h[0][-1][0],3) |
| def test_reward_requires_stop_and_charges_proposals(self): |
| binary=b''.join(w.to_bytes(4,'little') for w in (rv32.addi('x10','x0',7),rv32.EBREAK)); cases=[([1,2],7)] |
| self.assertAlmostEqual(terminal_reward(binary,binary,cases,0,[(MAX_WINDOW,0)]),.98) |
| self.assertLess(terminal_reward(binary,binary,cases,0,[(0,0)]),0) |
| self.assertLess(terminal_reward(binary,binary,cases,1,[(MAX_WINDOW,0)]),0) |
| self.assertLess(terminal_reward(binary,binary,[([1,2],9)],0,[(MAX_WINDOW,0)]),0) |
| def test_early_none(self): |
| p=Sequence([MAX_WINDOW]);out,e,h,_=rollout(p,['test'],[b'1234']) |
| self.assertEqual(out,[b'1234']);self.assertEqual(e,[0]);self.assertEqual(len(h[0]),1) |
|
|
| if __name__=='__main__':unittest.main() |
|
|
