"""Native comparisons for the instruction forms added for allocation/initialization."""
import json
import struct
import subprocess
import sys
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'tools'))
from a64_program import ProgramCompiler
from a64_lift import UnsupportedInstruction
OUT=ROOT/'build/initialization'
OUT.mkdir(parents=True,exist_ok=True)
def pack(words): return struct.pack('<'+'I'*len(words),*words)
def run(*args): subprocess.run(list(map(str,args)),cwd=ROOT,check=True)

operations=[]
for width in (32,64):
    for shift in range(width):
        operations.append((0x13000000 if width==32 else 0x93400000) | shift<<16 | (width-1)<<10)
    for length in range(1,6 if width==32 else 7):
        for operation in range(3):
            # Rotate repeated patterns at every supported element size.
            n=int(length==6)
            s=((((~((1 << (length+1))-1))&63) | ((1 << length)//2-1)))
            rotation=(1 << length)-1
            operations.append((0x12000000 if width==32 else 0x92000000) | operation<<29 | n<<22 | rotation<<16 | s<<10)
words=[]
for word in operations:
    words += [0xaa0203e0,word,0xf8008420] # mov x0,x2; op; str x0,[x1],#8
words += [0xd65f03c0]
compiler=ProgramCompiler(memory_pages=2)
(OUT/'initialization.wat').write_text(compiler.compile({0x1000:pack(words)},0x1000))
run(ROOT/'build/wabt/wat2wasm',OUT/'initialization.wat','-o',OUT/'initialization.wasm')
run(ROOT/'build/wabt/wasm2c',OUT/'initialization.wasm','-o',OUT/'initialization.c','--module-name=initialization')
(OUT/'oracle.S').write_text('.text\n.p2align 2\n.globl _native_initialization\n_native_initialization:\n'+
                           '\n'.join(f'.long 0x{w:08x}' for w in words)+'\n')
invalid=[
    [0x3c800000,0xd65f03c0], # uninitialized vector register
    [0xb4000040,0x6f00e400,0x3c800000,0xd65f03c0], # zero is missing on one incoming path
    [0x14000002,0x6f00e400,0x3c800000,0xd65f03c0], # branch into the idiom
    [0x6f00e400,0x3c800001,0xd65f03c0], # mismatched vector register
    [0x90000000,0xd65f03c0], # undeclared static page
    [0x9240fc00,0xd65f03c0], # reserved all-ones logical mask
    [0x93400000,0xd65f03c0], # signed bitfield beyond ASR
]
for fixture in invalid:
    try: ProgramCompiler().compile({0x1000:pack(fixture)},0x1000)
    except UnsupportedInstruction: pass
    else: raise AssertionError('Invalid initialization fixture accepted')
for options in [dict(memory_pages=0),dict(memory_pages=32769),dict(returning_hosts={1}),
                dict(static_pages={0:1}),dict(hosts={1:'same',2:'same'},returning_hosts={1})]:
    try: ProgramCompiler(**options)
    except ValueError: pass
    else: raise AssertionError('Invalid program configuration accepted')
(OUT/'verify.c').write_text(f'''
#include "initialization.h"
#include "wasm-rt-impl.h"
#include "wasm-rt-exceptions.h"
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
extern uint64_t native_initialization(uint64_t,uint64_t *,uint64_t);
int main(void) {{
 wasm_rt_init();w2c_initialization module;wasm2c_initialization_instantiate(&module);
 uint8_t *memory=w2c_initialization_memory(&module)->data;
 uint64_t seed=0xdeadbeef12345678,reference[{len(operations)}];
 for(unsigned i=0;i<512;++i) {{
  seed^=seed<<13;seed^=seed>>7;seed^=seed<<17;
  uint64_t input=i==0 ? 0 : i==1 ? UINT64_MAX : seed;
  uint64_t registers[32]={{0}};registers[1]=66048;registers[2]=input;
  memset(memory,0xa5,131072);memcpy(memory,registers,sizeof(registers));
  native_initialization(0,reference,input);
  assert(w2c_initialization_run(&module,0,10000)==0);
  assert(memcmp(memory+66048,reference,sizeof(reference))==0);
  for(unsigned j=256;j<131072;++j) {{
   if(j>=66048 && j<66048+sizeof(reference)) continue;
   assert(memory[j]==0xa5);
  }}
 }}
 const uint64_t bad[]={{131072-7,131072,UINT64_C(0x100000000)+66048,UINT64_MAX-7}};
 for(unsigned i=0;i<sizeof(bad)/sizeof(bad[0]);++i) {{
  uint64_t registers[32]={{0}};registers[1]=bad[i];
  memset(memory,0xa5,131072);memcpy(memory,registers,sizeof(registers));
  wasm_rt_trap_t trap=wasm_rt_impl_try();
  if(!trap) {{w2c_initialization_run(&module,0,10000);assert(0);}}
  assert(trap==WASM_RT_TRAP_UNREACHABLE);
  for(unsigned j=256;j<131072;++j) assert(memory[j]==0xa5);
 }}
 wasm2c_initialization_free(&module);wasm_rt_free();
 puts("PASS: {len(operations)*512} native ASR/logical-mask comparisons, extended memory, 64-bit address guards, 12 rejection cases");
}}
''')
runtime=ROOT/'third_party/wabt/wasm2c'
run('xcrun','clang','-O2','-DWASM_RT_USE_MMAP=0','-DWASM_RT_MEMCHECK_BOUNDS_CHECK=1',
    '-I'+str(OUT),'-I'+str(runtime),OUT/'initialization.c',OUT/'verify.c',OUT/'oracle.S',
    runtime/'wasm-rt-impl.c',runtime/'wasm-rt-mem-impl.c',runtime/'wasm-rt-exceptions-impl.c','-o',OUT/'verify')
run(OUT/'verify')
(OUT/'verification.json').write_text(json.dumps({'status':'pass','native_comparisons':len(operations)*512,
  'instruction_forms':len(operations),'invalid_cases':12,'source':'hand-authored instruction fixtures'},indent=2)+'\n')