from pathlib import Path
from quantum_ram import apply_single_gate
ROOT=Path(__file__).resolve().parents[1]
P=2**64-2**32+1
G=[(2,3),(5,7),(11,13),(17,19)]
def cmul(a,b):return ((a[0]*b[0]-a[1]*b[1])%P,(a[0]*b[1]+a[1]*b[0])%P)
def cadd(a,b):return ((a[0]+b[0])%P,(a[1]+b[1])%P)
def enc(xs):return '['+', '.join('"'+str(x)+'"' for x in xs)+']'
(ROOT/'rs/tests/fixtures/quantum_ram.tasm').write_text('// Generated by scripts/quantum_ram.py; benchmark-private scratch [2^40, 2^40+64).\n'+apply_single_gate())
for n in range(1,4):
size=1<<n;original=[(3*(2*i+1),3*(2*i+2)) for i in range(size)]
cells={99:987,100+2*size:654}
cells.update({100+i:v for i,v in enumerate(x for pair in original for x in pair)})
addresses=sorted(cells)
init='\n'.join(f'mem.write({a},{v})' for a,v in sorted(cells.items()))
gateargs=','.join(f'gates.Complex{{re:{a},im:{b}}}' for a,b in G)
source='program quantum_ram\nuse std.quantum.gates\nuse vm.io.mem\nfn main(){\n'+init+'\nlet addr:Field=pub_read()\nlet n:Field=pub_read()\nlet target:Field=pub_read()\ngates.apply_single_gate(addr,n,target,'+gateargs+')\n'+'\n'.join(f'pub_write(mem.read({a}))' for a in addresses)+'\n}\n'
prefix=' '.join(f'push {v} push {a} write_mem 1 pop 1' for a,v in sorted(cells.items()))+' read_io 1 read_io 1 read_io 1 '+' '.join(f'push {v}' for pair in G for v in pair)+' call __apply_single_gate '+' '.join(f'push {a} read_mem 1 pop 1 write_io 1' for a in addresses)+' halt'
d=ROOT/f'baselines/triton/fixtures/quantum-ram-{n}';d.mkdir(exist_ok=True);(d/'main.tri').write_text(source)
for target in range(n):
result=list(original)
for i in range(size):
if i & (1<<target):continue
j=i+(1<<target);a=original[i];b=original[j]
result[i]=cadd(cmul(G[0],a),cmul(G[1],b));result[j]=cadd(cmul(G[2],a),cmul(G[3],b))
output=[987]+[v for pair in result for v in pair]+[654]
fixture='source = "main.tri"\nhand = "../../std/quantum/gates.tasm"\ninput = '+enc([100,n,target])+'\noutput = '+enc(output)+'\nreference = "Independent integer complex 2x2 matrix arithmetic modulo Goldilocks, every target n=1..3; neighboring cells unchanged. Hand scratch [2^40,2^40+64) is benchmark-private and clobbered."\nhand_prefix = "'+prefix+'"\n'
(d/f'target-{target}.bench.toml').write_text(fixture)
if n==1:
for name,args in [('too-many',[100,13,0]),('target',[100,1,1]),('address',[(1<<32)-3,1,0])]:
(d/f'invalid-{name}.bench.toml').write_text(fixture.replace('input = '+enc([100,n,0]),'input = '+enc(args)).replace('output = '+enc(output),'output = []\nexpect_failure = true'))