module Data.SHA256Schedule

import Data.SHA256
import Model.Config
import Std.Natural

family SHA256WordReadResult : Type 0
constructor SHA256WordReadSucceeded
field unrestricted sha256WordReadValue : (family ModelWord32)
field unrestricted sha256WordReadRemaining : Bytes
constructor SHA256WordReadFailed
field unrestricted sha256WordReadError : (family SHA256ErrorCode)

end-family

family SHA256BlockDecodeResult : Type 0
constructor SHA256BlockDecodeSucceeded
field unrestricted sha256BlockDecodedSchedule : (family SHA256Schedule)
field unrestricted sha256BlockDecodedWordCount : Nat
constructor SHA256BlockDecodeFailed
field unrestricted sha256BlockDecodeError : (family SHA256ErrorCode)
field unrestricted sha256BlockDecodeWordOrdinal : Nat

end-family

family SHA256ScheduleLookupResult : Type 0
constructor SHA256ScheduleLookupSucceeded
field unrestricted sha256ScheduleLookupWord : (family ModelWord32)
constructor SHA256ScheduleLookupFailed
field unrestricted sha256ScheduleLookupError : (family SHA256ErrorCode)
field unrestricted sha256ScheduleLookupIndex : Nat

end-family

def sha256NaturalFour =
  (byte-to-nat (byte 4))

def sha256NaturalSixteen =
  (byte-to-nat (byte 16))

def sha256NaturalSixtyFour =
  (byte-to-nat (byte 64))

def sha256ReadWord =
  (lambda unrestricted input : Bytes .
    (nat-eliminate
      (lambda unrestricted sufficient : Nat . (family SHA256WordReadResult))
      (constructor
        SHA256WordReadResult
        SHA256WordReadFailed
        (constructor SHA256ErrorCode SHA256BlockLengthInvalid))
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : (family SHA256WordReadResult) .
          (app
            (lambda unrestricted tail1 : Bytes .
              (app
                (lambda unrestricted tail2 : Bytes .
                  (app
                    (lambda unrestricted tail3 : Bytes .
                      (app
                        (lambda unrestricted tail4 : Bytes .
                          (constructor
                            SHA256WordReadResult
                            SHA256WordReadSucceeded
                            (constructor
                              ModelWord32
                              ModelWord32Value
                              (bytes-head tail3)
                              (bytes-head tail2)
                              (bytes-head tail1)
                              (bytes-head input))
                            tail4))
                        (bytes-tail tail3)))
                    (bytes-tail tail2)))
                (bytes-tail tail1)))
            (bytes-tail input))))
      (naturalLessOrEqual sha256NaturalFour (bytes-length input))))

def sha256DecodeBlockWordsWithFuel =
  (lambda unrestricted fuel : Nat .
    (nat-eliminate
      (lambda unrestricted current : Nat .
        (pi unrestricted input : Bytes .
          (pi unrestricted ordinal : Nat . (family SHA256BlockDecodeResult))))
      (lambda unrestricted input : Bytes .
        (lambda unrestricted ordinal : Nat .
          (nat-eliminate
            (lambda unrestricted emptyFlag : Nat . (family SHA256BlockDecodeResult))
            (constructor
              SHA256BlockDecodeResult
              SHA256BlockDecodeFailed
              (constructor SHA256ErrorCode SHA256BlockLengthInvalid)
              ordinal)
            (lambda unrestricted predecessor : Nat .
              (lambda unrestricted induction : (family SHA256BlockDecodeResult) .
                (constructor
                  SHA256BlockDecodeResult
                  SHA256BlockDecodeSucceeded
                  (constructor SHA256Schedule SHA256ScheduleEnd)
                  ordinal)))
            (naturalIsZero (bytes-length input)))))
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : (pi unrestricted input : Bytes . (pi unrestricted ordinal : Nat . (family SHA256BlockDecodeResult))) .
          (lambda unrestricted input : Bytes .
            (lambda unrestricted ordinal : Nat .
              (eliminate
                SHA256WordReadResult
                (lambda unrestricted current : (family SHA256WordReadResult) .
                  (family SHA256BlockDecodeResult))
                (sha256ReadWord input)
                (branch
                  SHA256WordReadSucceeded
                  word
                  remaining
                  .
                  (eliminate
                    SHA256BlockDecodeResult
                    (lambda unrestricted current : (family SHA256BlockDecodeResult) .
                      (family SHA256BlockDecodeResult))
                    (induction remaining (succ ordinal))
                    (branch
                      SHA256BlockDecodeSucceeded
                      tailSchedule
                      finalCount
                      .
                      (constructor
                        SHA256BlockDecodeResult
                        SHA256BlockDecodeSucceeded
                        (constructor SHA256Schedule SHA256ScheduleNext word tailSchedule)
                        finalCount))
                    (branch
                      SHA256BlockDecodeFailed
                      error
                      failedOrdinal
                      .
                      (constructor
                        SHA256BlockDecodeResult
                        SHA256BlockDecodeFailed
                        error
                        failedOrdinal))))
                (branch
                  SHA256WordReadFailed
                  error
                  .
                  (constructor SHA256BlockDecodeResult SHA256BlockDecodeFailed error ordinal)))))))
      fuel))

def sha256DecodeBlockWords =
  (lambda unrestricted block : Bytes .
    (nat-eliminate
      (lambda unrestricted validLength : Nat . (family SHA256BlockDecodeResult))
      (constructor
        SHA256BlockDecodeResult
        SHA256BlockDecodeFailed
        (constructor SHA256ErrorCode SHA256BlockLengthInvalid)
        zero)
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : (family SHA256BlockDecodeResult) .
          (sha256DecodeBlockWordsWithFuel sha256NaturalSixteen block zero)))
      (naturalEqual (bytes-length block) sha256NaturalSixtyFour)))

-- Drop the first `index` nodes of a schedule, returning the suffix that begins
-- at that index (or the empty schedule when the index runs past the end).
--
-- This is a FIRST-ORDER fold: the accumulator is a `SHA256Schedule` VALUE, not a
-- function. At each step it destructures the current suffix and keeps the tail --
-- a shared sub-structure of the original list, no copy -- so producing the
-- index-th suffix costs O(index) constant-work steps. (The list-recursion
-- hypothesis `nodeInduction` is deliberately unused, so it is never materialized.)
def sha256ScheduleDrop =
  (lambda unrestricted index : Nat .
    (lambda unrestricted schedule : (family SHA256Schedule) .
      (nat-eliminate
        (lambda unrestricted current : Nat . (family SHA256Schedule))
        schedule
        (lambda unrestricted predecessor : Nat .
          (lambda unrestricted induction : (family SHA256Schedule) .
            (eliminate
              SHA256Schedule
              (lambda unrestricted current : (family SHA256Schedule) . (family SHA256Schedule))
              induction
              (branch SHA256ScheduleEnd . (constructor SHA256Schedule SHA256ScheduleEnd))
              (branch SHA256ScheduleNext word tail nodeInduction . tail))))
        index)))

-- Linear O(index) message-schedule lookup.
--
-- The previous definition walked the linked list with a `nat-eliminate` over the
-- FULL index at every node, and its successor step re-invoked the list-recursion
-- hypothesis (`(app induction predecessor)`) at every intermediate fold level.
-- Because the VM evaluates a fold's induction eagerly, one lookup at index k
-- forced a fresh lookup of the tail at 0,1,...,k-1 -- an exponential re-walk that
-- made schedule expansion run in ~O(N^4) (~3.2 billion evals for a single 4-byte
-- hash). A HIGHER-ORDER rewrite (function-valued accumulator) removes the
-- re-invocation but still pays O(k^2) per lookup, because the VM symbolically
-- unrolls the function accumulator into a term of size O(k) at every step.
--
-- This version is FIRST-ORDER: `sha256ScheduleDrop` folds the SCHEDULE VALUE
-- itself to the suffix at `index` (O(index), shared sub-structures, no term
-- building), then reads that suffix's head word. Result values are byte-for-byte
-- identical to the original (index i still yields W[i]); only the cost changed,
-- so the SHA-256 digest is preserved exactly.
def sha256ScheduleLookup =
  (lambda unrestricted schedule : (family SHA256Schedule) .
    (lambda unrestricted index : Nat .
      (eliminate
        SHA256Schedule
        (lambda unrestricted current : (family SHA256Schedule) .
          (family SHA256ScheduleLookupResult))
        (sha256ScheduleDrop index schedule)
        (branch
          SHA256ScheduleEnd
          .
          (constructor
            SHA256ScheduleLookupResult
            SHA256ScheduleLookupFailed
            (constructor SHA256ErrorCode SHA256ScheduleLengthInvalid)
            index))
        (branch
          SHA256ScheduleNext
          word
          tail
          nodeInduction
          .
          (constructor SHA256ScheduleLookupResult SHA256ScheduleLookupSucceeded word)))))

def sha256ScheduleAppend =
  (lambda unrestricted schedule : (family SHA256Schedule) .
    (lambda unrestricted word : (family ModelWord32) .
      (eliminate
        SHA256Schedule
        (lambda unrestricted current : (family SHA256Schedule) . (family SHA256Schedule))
        schedule
        (branch
          SHA256ScheduleEnd
          .
          (constructor
            SHA256Schedule
            SHA256ScheduleNext
            word
            (constructor SHA256Schedule SHA256ScheduleEnd)))
        (branch
          SHA256ScheduleNext
          head
          tail
          induction
          .
          (constructor SHA256Schedule SHA256ScheduleNext head induction)))))

def sha256ScheduleLength =
  (lambda unrestricted schedule : (family SHA256Schedule) .
    (eliminate
      SHA256Schedule
      (lambda unrestricted current : (family SHA256Schedule) . Nat)
      schedule
      (branch SHA256ScheduleEnd . zero)
      (branch SHA256ScheduleNext word tail induction . (succ induction))))