module Compiler.ELF

import Compiler.AST
import Compiler.Lexer
import Compiler.Parser
import Compiler.Elaborator
import Compiler.Codegen
import Model.Parameter
import Model.Word64
import Std.Natural

family ExecutableResult : Type 0
constructor ExecutableGenerated
field unrestricted executableBytes : Bytes
constructor ExecutableUnboundVariable
field unrestricted executableUnboundSpelling : Bytes
constructor ExecutableUnsupportedTerm
field unrestricted executableUnsupportedCode : Nat
constructor ExecutableUnsupportedSuccessor

end-family

-- The ELF machines the physical executor is built for: e_machine 62
-- (EM_X86_64) or 183 (EM_AARCH64).  Everything else in the header -- entry
-- 0x400078, one PT_LOAD at 0x400000 with a 4 MiB memory extent, 4 KiB
-- alignment -- is the same for both.
def elfMachineX86_64 : Byte = (byte 62)
def elfMachineAArch64 : Byte = (byte 183)

-- the ELF identification and e_type, before e_machine
def elfHeaderBeforeMachine : Bytes =
  (bytes 127 69 76 70 2 1 1 0 0 0 0 0 0 0 0 0 2 0)

-- from e_machine's second byte to the program header's p_filesz
def elfHeaderAfterMachine : Bytes =
  (bytes 0 1 0 0 0 120 0 64 0 0 0 0 0 64 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 64 0 56 0 1 0 0 0 0 0 0 0 1 0 0 0 5 0 0 0 0 0 0 0 0 0 0 0 0 0 64 0 0 0 0 0 0 0 64 0 0 0 0 0)

def elfHeaderPrefix : Bytes =
  (bytes-append elfHeaderBeforeMachine (bytes-append (bytes elfMachineX86_64) elfHeaderAfterMachine))

def elfLoadSize64k : Bytes =
  (bytes 0 0 4 0 0 0 0 0)

def elfHeaderAlign : Bytes =
  (bytes 0 16 0 0 0 0 0 0)

def fixedZeroELFHeader : Bytes =
  (bytes-append
    elfHeaderPrefix
    (bytes-append elfLoadSize64k (bytes-append elfLoadSize64k elfHeaderAlign)))

def zeroPadding1 : Bytes =
  (bytes 0)

def zeroPadding2 : Bytes =
  (bytes-append zeroPadding1 zeroPadding1)

def zeroPadding4 : Bytes =
  (bytes-append zeroPadding2 zeroPadding2)

def zeroPadding8 : Bytes =
  (bytes-append zeroPadding4 zeroPadding4)

def zeroPadding16 : Bytes =
  (bytes-append zeroPadding8 zeroPadding8)

def zeroPadding32 : Bytes =
  (bytes-append zeroPadding16 zeroPadding16)

def zeroPadding64 : Bytes =
  (bytes-append zeroPadding32 zeroPadding32)

def zeroPadding128 : Bytes =
  (bytes-append zeroPadding64 zeroPadding64)

def zeroPadding256 : Bytes =
  (bytes-append zeroPadding128 zeroPadding128)

def zeroPadding512 : Bytes =
  (bytes-append zeroPadding256 zeroPadding256)

def zeroPadding1024 : Bytes =
  (bytes-append zeroPadding512 zeroPadding512)

def zeroPadding2048 : Bytes =
  (bytes-append zeroPadding1024 zeroPadding1024)

def zeroPadding4096 : Bytes =
  (bytes-append zeroPadding2048 zeroPadding2048)

def zeroPadding8192 : Bytes =
  (bytes-append zeroPadding4096 zeroPadding4096)

def zeroPadding16384 : Bytes =
  (bytes-append zeroPadding8192 zeroPadding8192)

def zeroPadding32768 : Bytes =
  (bytes-append zeroPadding16384 zeroPadding16384)

def zeroPadding65536 : Bytes =
  (bytes-append zeroPadding32768 zeroPadding32768)

def zeroPadding131072 : Bytes =
  (bytes-append zeroPadding65536 zeroPadding65536)

def zeroPadding262144 : Bytes =
  (bytes-append zeroPadding131072 zeroPadding131072)

def wrapMachineCode =
  (lambda unrestricted machineCode : Bytes .
    (bytes-append (bytes-append fixedZeroELFHeader machineCode) zeroPadding262144))

-- Larger LOAD segment (4 MiB) for the PHYSICAL native-program path only (nvctl/GPU probes whose
-- interpreted image body + embedded kernel + upload payloads exceed the 256 KiB fixedZeroELFHeader
-- segment). The header LENGTH is byte-identical to fixedZeroELFHeader (120 bytes), so the machine
-- code still starts at file offset 120 and the interpreter's state/scratch base is unchanged; only
-- the mapped segment END moves out from 0x440000 to 0x800000. The self-host compiler ELF path
-- (compileClosedZeroExecutable -> wrapMachineCode) is deliberately NOT changed, so its byte-identity
-- is preserved. p_filesz = p_memsz = 0x00400000 (4 MiB); the file is padded to cover it.  The old
-- 1-MiB physical segment was smaller than a valid 4,242-command checkpoint-resume program: Linux
-- mapped only 0x400000..0x500000 and execution faulted exactly at the unmapped 0x500000 boundary.
def zeroPadding524288 : Bytes =
  (bytes-append zeroPadding262144 zeroPadding262144)

def zeroPadding1048576 : Bytes =
  (bytes-append zeroPadding524288 zeroPadding524288)

def zeroPadding2097152 : Bytes =
  (bytes-append zeroPadding1048576 zeroPadding1048576)

def zeroPadding4194304 : Bytes =
  (bytes-append zeroPadding2097152 zeroPadding2097152)

def elfLoadSize4M : Bytes =
  (bytes 0 0 64 0 0 0 0 0)

def fixedZeroELFHeaderLarge : Bytes =
  (bytes-append
    elfHeaderPrefix
    (bytes-append elfLoadSize4M (bytes-append elfLoadSize4M elfHeaderAlign)))

def wrapMachineCodeLarge =
  (lambda unrestricted machineCode : Bytes .
    (bytes-append (bytes-append fixedZeroELFHeaderLarge machineCode) zeroPadding4194304))

-- Encode a model word in the little-endian byte order used by ELF64 fields.
def elfModelWord64Bytes =
  (lambda unrestricted value : (family ModelWord64) .
    (eliminate
      ModelWord64
      (lambda unrestricted current : (family ModelWord64) . Bytes)
      value
      (branch ModelWord64Value b0 b1 b2 b3 b4 b5 b6 b7 .
        (bytes b0 b1 b2 b3 b4 b5 b6 b7))))

-- The physical executor still receives the same 4 MiB PT_LOAD memory extent,
-- but the file carries only its header and machine bytes, padded with zeros
-- to the next 0x1000-byte page.  Linux zero-fills p_memsz - p_filesz,
-- avoiding a 4 MiB compile-time Bytes value.  The page padding is not
-- optional: the segment is not writable, and a kernel before 6.7 zeroes the
-- rest of the last file-backed page itself (padzero) and fails the exec with
-- EFAULT when it cannot write it -- measured 2026-09-24 on a RunPod host with
-- Linux 6.5.0-44: every artifact died in execve with SIGSEGV; the same
-- bytes padded to a page loaded.  With p_filesz ending on a page boundary
-- there is no partial page to zero.
def elfPageBytes : Nat = 4096

-- zeroPadding<bit> when that bit of n is set
def elfZerosIfBit =
  (lambda unrestricted n : Nat . (lambda unrestricted bit : Nat . (lambda unrestricted chunk : Bytes .
    (nat-eliminate (lambda unrestricted current : Nat . Bytes) b""
      (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : Bytes . chunk))
      (naturalModuloUnchecked (naturalDivideUnchecked n bit) 2)))))

-- n zero bytes, n below a page
def elfZerosBelowPage =
  (lambda unrestricted n : Nat .
    (bytes-append (elfZerosIfBit n 2048 zeroPadding2048) (bytes-append (elfZerosIfBit n 1024 zeroPadding1024)
    (bytes-append (elfZerosIfBit n 512 zeroPadding512) (bytes-append (elfZerosIfBit n 256 zeroPadding256)
    (bytes-append (elfZerosIfBit n 128 zeroPadding128) (bytes-append (elfZerosIfBit n 64 zeroPadding64)
    (bytes-append (elfZerosIfBit n 32 zeroPadding32) (bytes-append (elfZerosIfBit n 16 zeroPadding16)
    (bytes-append (elfZerosIfBit n 8 zeroPadding8) (bytes-append (elfZerosIfBit n 4 zeroPadding4)
    (bytes-append (elfZerosIfBit n 2 zeroPadding2) (elfZerosIfBit n 1 zeroPadding1)))))))))))))

-- the zeros that carry header + machine code to the next page
def elfPagePadding =
  (lambda unrestricted machineCode : Bytes .
    (elfZerosBelowPage
      (naturalModuloUnchecked
        (naturalSaturatingSubtract elfPageBytes
          (naturalModuloUnchecked (naturalAdd (byte-to-nat (byte 120)) (bytes-length machineCode)) elfPageBytes))
        elfPageBytes)))

-- elfHeaderPrefix for another e_machine
def elfHeaderPrefixFor =
  (lambda unrestricted machine : Byte .
    (bytes-append elfHeaderBeforeMachine (bytes-append (bytes machine) elfHeaderAfterMachine)))

def fixedZeroELFHeaderSparseLargeFor =
  (lambda unrestricted machine : Byte .
    (lambda unrestricted machineCode : Bytes .
      (let unrestricted fileExtent =
        (modelWord64FromNaturalTruncated
          (naturalAdd (byte-to-nat (byte 120)) (naturalAdd (bytes-length machineCode) (bytes-length (elfPagePadding machineCode)))))
      in
        (bytes-append
          (elfHeaderPrefixFor machine)
          (bytes-append
            (elfModelWord64Bytes fileExtent)
            (bytes-append elfLoadSize4M elfHeaderAlign))))))

def wrapMachineCodeSparseLargeFor =
  (lambda unrestricted machine : Byte .
    (lambda unrestricted machineCode : Bytes .
      (bytes-append (fixedZeroELFHeaderSparseLargeFor machine machineCode) (bytes-append machineCode (elfPagePadding machineCode)))))

-- The same container with its memory extent exactly its file extent.  The
-- 4 MiB extent above dates from when the file itself had to cover the
-- image; the file now does, and nothing is read from memory beyond it (an
-- embedded artifact's payloads are read with pread).  The AArch64 host uses
-- this one: a loader that maps no zero-filled tail is also one qemu-user
-- accepts without the segment being writable, so the executable its tests
-- run is the executable itself.  x86-64 keeps the 4 MiB extent, which is
-- what its published executables carry.
def fixedZeroELFHeaderExactFor =
  (lambda unrestricted machine : Byte .
    (lambda unrestricted machineCode : Bytes .
      (let unrestricted fileExtent =
        (elfModelWord64Bytes
          (modelWord64FromNaturalTruncated
            (naturalAdd (byte-to-nat (byte 120)) (naturalAdd (bytes-length machineCode) (bytes-length (elfPagePadding machineCode))))))
      in
        (bytes-append
          (elfHeaderPrefixFor machine)
          (bytes-append fileExtent (bytes-append fileExtent elfHeaderAlign))))))

def wrapMachineCodeExactFor =
  (lambda unrestricted machine : Byte .
    (lambda unrestricted machineCode : Bytes .
      (bytes-append (fixedZeroELFHeaderExactFor machine machineCode) (bytes-append machineCode (elfPagePadding machineCode)))))

def fixedZeroELFHeaderSparseLarge = (fixedZeroELFHeaderSparseLargeFor elfMachineX86_64)

def wrapMachineCodeSparseLarge = (wrapMachineCodeSparseLargeFor elfMachineX86_64)

def fixedZeroELF : Bytes =
  (wrapMachineCode fixedZeroMachineCode)

def compileClosedZeroExecutable =
  (lambda unrestricted result : (family ClosedNaturalElaboration) .
    (eliminate
      CodegenResult
      (lambda unrestricted value : (family CodegenResult) . (family ExecutableResult))
      (compileClosedNaturalElaboration result)
      (branch
        CodeGenerated
        machineCode
        .
        (constructor ExecutableResult ExecutableGenerated (wrapMachineCode machineCode)))
      (branch
        CodegenUnboundVariable
        codegenUnboundSpelling
        .
        (constructor ExecutableResult ExecutableUnboundVariable codegenUnboundSpelling))
      (branch
        CodegenUnsupportedTerm
        codegenUnsupportedCode
        .
        (constructor ExecutableResult ExecutableUnsupportedTerm codegenUnsupportedCode))))

def executableSample : (family ExecutableResult) =
  (compileClosedZeroExecutable elaboratedParserSample)

def executableSize : Nat =
  (eliminate
    ExecutableResult
    (lambda unrestricted result : (family ExecutableResult) . Nat)
    executableSample
    (branch ExecutableGenerated executableBytes . (bytes-length executableBytes))
    (branch ExecutableUnboundVariable executableUnboundSpelling . zero)
    (branch ExecutableUnsupportedTerm executableUnsupportedCode . zero)
    (branch ExecutableUnsupportedSuccessor . zero))