module Runtime.NativeIndependentTree

import Learning.IndependentState
import Component.Pointwise
import Checkpoint.Envelope
import Checkpoint.NativeEnvelope
import Checkpoint.NativeRegions
import Checkpoint.NativePeerPublication
import Checkpoint.NativeTreePublication
import Checkpoint.RegionManifest
import Compiler.Planning.RegionFrame
import Representation.ParticipantTree
import Representation.Partition
import Runtime.NativeIndependentState
import Runtime.NativePointwisePlan
import Runtime.NativePhysicalProgram
import Runtime.NativePhysicalProgramComposer
import Runtime.NativeHostTarget
import Runtime.NativeResourcePlan
import Runtime.NativeLaunchRecipeLoad
import Runtime.NativeMemoryPlan
import Runtime.NativeSocketPlan
import Runtime.NativePeerEndpoint
import Runtime.NativePeerTree
import Runtime.NativeRegionFrame
import Runtime.NativeWordArithmetic
import Runtime.RegionTransfer
import Runtime.ArenaCertificate
import Learning.Transition
import Data.Bytes
import Model.Word64
import Std.Foundation
import Std.List
import Std.Natural
import Std.Word

-- The independent-state learner over a bounded-fan-out hierarchy instead of a
-- star: owners 0..n-1 sit at the roles of Representation.ParticipantTree,
-- each establishes only its parent and child edges
-- (Runtime.NativePeerTree), trains its own state, and publishes through
-- Checkpoint.NativeTreePublication: own shard and group manifest locally,
-- one commitment up, release down. No owner receives another owner's state,
-- and an owner's arena and edges are bounded by the fan-out and height, not
-- by n. Resume reads only the root-to-owner manifest path and the owner's
-- own shard. The local equation, numerics and checkpoint identity of the
-- state itself are the star realization's (Runtime.NativeIndependentState).
family NativeIndependentTreeLayout : Type 0
constructor NativeIndependentTreeLayoutOf
-- template frames, planned for the edge of role 1; every edge's frames share
-- this storage (Runtime.NativeIndependentState.niFrameAt)
field unrestricted independentTreeCommitment : (family RegionFrameLayout)
field unrestricted independentTreeRelease : (family RegionFrameLayout)
field unrestricted independentTreeCheckpoint : (family NativeTreeCheckpoint)
field unrestricted independentTreeScratchWords : Nat
-- the system identity's digest, computed once per artifact
field unrestricted independentTreeDigest : Bytes

end-family

def nitCommitment =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (eliminate
      NativeIndependentTreeLayout
      (lambda unrestricted current : (family NativeIndependentTreeLayout) .
        (family RegionFrameLayout))
      layout
      (branch
        NativeIndependentTreeLayoutOf
        commitment
        release
        checkpoint
        scratch
        digest
        .
        commitment)))

def nitRelease =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (eliminate
      NativeIndependentTreeLayout
      (lambda unrestricted current : (family NativeIndependentTreeLayout) .
        (family RegionFrameLayout))
      layout
      (branch NativeIndependentTreeLayoutOf commitment release checkpoint scratch digest . release)))

def nitCheckpoint =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (eliminate
      NativeIndependentTreeLayout
      (lambda unrestricted current : (family NativeIndependentTreeLayout) .
        (family NativeTreeCheckpoint))
      layout
      (branch
        NativeIndependentTreeLayoutOf
        commitment
        release
        checkpoint
        scratch
        digest
        .
        checkpoint)))

def nitScratchWords =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (eliminate
      NativeIndependentTreeLayout
      (lambda unrestricted current : (family NativeIndependentTreeLayout) . Nat)
      layout
      (branch NativeIndependentTreeLayoutOf commitment release checkpoint scratch digest . scratch)))

def nitDigest =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (eliminate
      NativeIndependentTreeLayout
      (lambda unrestricted current : (family NativeIndependentTreeLayout) . Bytes)
      layout
      (branch NativeIndependentTreeLayoutOf commitment release checkpoint scratch digest . digest)))

def nitStateAt =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (nativeCheckpointAlign (nativeTreeEnd (nitCheckpoint layout))))

def nitValuesAt =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (naturalAdd (nitStateAt layout) nativeRegionCaptureBytes))

def nitEndpoint =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
      (nativePointwiseAt (nitValuesAt layout) (naturalAdd slots (nitScratchWords layout)))))

def nitExtent =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
      (lambda unrestricted edges : (family StdList (family NativePeerEdge)) .
        (nativePeerEdgesEnd (nitEndpoint slots layout) edges))))

-- A role holds one interval of the one logical region: its own state words.
def nitOwnOf =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted role : Nat .
        (constructor
          StdList
          StdListCons
          (family CheckpointRegionPartition)
          (constructor
            CheckpointRegionPartition
            CheckpointRegionPartitionOf
            0
            (constructor
              StdList
              StdListCons
              (family LogicalInterval)
              (niInterval slots system role)
              (constructor StdList StdListEmpty (family LogicalInterval))))
          (constructor StdList StdListEmpty (family CheckpointRegionPartition))))))

-- The tree protocol is a different publication from the star's, so its
-- checkpoint identity differs even for the same system and owners.
def nitContext =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (constructor
        CheckpointEnvelopeContract
        CheckpointEnvelopeContractValue
        (bytes-append
          b"independent-tree-state-f64-finite-rne-separate-v1;"
          (independentStateIdentity slots system))
        b"typed-sequential-local-transitions"
        b"immutable-local-inputs;one-input-round-per-update"
        b"declared-initial-state;no-randomness"
        1
        1
        1)))

def nitParent =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat .
      (lambda unrestricted role : Nat .
        (participantTreeParentOf (participantTreeNode fanout count role)))))

-- The frames of the edge between `child` and its parent.
def nitCommitmentFor =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted fanout : Nat .
        (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
          (lambda unrestricted child : Nat .
            (niFrameAt
              (niIdentity
                slots
                system
                (nitDigest layout)
                2
                child
                (nitParent fanout (niOwners slots system) child))
              nativeTreeCommitmentBytes
              (nitCommitment layout)))))))

def nitReleaseFor =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted fanout : Nat .
        (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
          (lambda unrestricted child : Nat .
            (niFrameAt
              (niIdentity
                slots
                system
                (nitDigest layout)
                3
                (nitParent fanout (niOwners slots system) child)
                child)
              nativePointwiseWordBytes
              (nitRelease layout)))))))

def nitWithLayout =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted fanout : Nat .
        (lambda unrestricted role : Nat .
          (lambda unrestricted epoch : Nat .
            (lambda unrestricted base : Nat .
              (lambda unrestricted scratchWords : Nat .
                (lambda unrestricted available : Nat .
                  (lambda unrestricted emit : (pi unrestricted layout : (family NativeIndependentTreeLayout) . Bytes) .
                    (let unrestricted digest =
                      (niIdentityDigest slots system)
                      in
                      (eliminate
                        StdResult
                        (lambda unrestricted result : (family StdResult (family RegionFrameLayoutFailure) (family RegionFrameLayout)) .
                          Bytes)
                        (planRegionFrameLayout
                          (niIdentity slots system digest 2 1 0)
                          nativeTreeCommitmentBytes
                          512
                          base
                          available)
                        (branch StdFailure error . b"")
                        (branch
                          StdSuccess
                          commitment
                          .
                          (eliminate
                            StdResult
                            (lambda unrestricted result : (family StdResult (family RegionFrameLayoutFailure) (family RegionFrameLayout)) .
                              Bytes)
                            (planRegionFrameLayout
                              (niIdentity slots system digest 3 0 1)
                              nativePointwiseWordBytes
                              512
                              (frameAlign (frameEnd commitment))
                              available)
                            (branch StdFailure error . b"")
                            (branch
                              StdSuccess
                              release
                              .
                              (eliminate
                                StdResult
                                (lambda unrestricted result : (family StdResult (family NativeCheckpointFailure) (family NativeTreeCheckpoint)) .
                                  Bytes)
                                (planNativeTreeCheckpoint
                                  (nitContext slots system)
                                  (niRegions slots system)
                                  fanout
                                  (niOwners slots system)
                                  role
                                  (nitOwnOf slots system)
                                  epoch
                                  (frameAlign (frameEnd release))
                                  available)
                                (branch StdFailure error . b"")
                                (branch
                                  StdSuccess
                                  checkpoint
                                  .
                                  (emit
                                    (constructor
                                      NativeIndependentTreeLayout
                                      NativeIndependentTreeLayoutOf
                                      commitment
                                      release
                                      checkpoint
                                      scratchWords
                                      digest))))))))))))))))))

-- argv for n owners (positions after the program name): 1 the previous root
-- manifest; 1 + r the previous group manifest of role r >= 1; n + 1 + r the
-- previous shard of role r; 2n + 1 + 2r and 2n + 2 + 2r the candidate and
-- committed group manifest of role r (role 0's is the published root);
-- 4n + 1 + 2r and 4n + 2 + 2r role r's candidate and committed shard; 6n + 1
-- the generation directory. Every owner receives the same argv and reads
-- only its own entries and its ancestors'.
def nitPaths =
  (lambda unrestricted owners : Nat .
    (lambda unrestricted role : Nat .
      (let unrestricted groups =
        (succ (naturalMultiply 2 owners))
        in
        (let unrestricted shards =
          (succ (naturalMultiply 4 owners))
          in
          (constructor
            NativeTreePaths
            NativeTreePathsOf
            (lambda unrestricted at : Nat . (succ at))
            (lambda unrestricted index : Nat . (succ (naturalAdd owners role)))
            (naturalAdd groups (naturalMultiply 2 role))
            (succ (naturalAdd groups (naturalMultiply 2 role)))
            (lambda unrestricted index : Nat . (naturalAdd shards (naturalMultiply 2 role)))
            (lambda unrestricted index : Nat . (succ (naturalAdd shards (naturalMultiply 2 role))))
            (succ (naturalMultiply 6 owners)))))))

-- Fresh: version, cursor 0 and the plan's epoch. Resume: the verified path
-- restore, the restored prefix, then epoch + 1 (a new membership epoch).
def nitStart =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
        (lambda unrestricted fanout : Nat .
          (lambda unrestricted role : Nat .
            (lambda unrestricted resume : Nat .
              (lambda unrestricted tail : (family NativePhysicalCommands) .
                (let unrestricted stateAt =
                  (nitStateAt layout)
                  in
                  (nat-eliminate
                    (lambda unrestricted mode : Nat . (family NativePhysicalCommands))
                    (nativeMemoryStoreWord64
                      b"independent-version"
                      stateAt
                      (nativeLaunchRecipeImmediate 0)
                      (nativeMemoryStoreWord64
                        b"independent-cursor"
                        (nativePointwiseAt stateAt 1)
                        (nativeLaunchRecipeImmediate 0)
                        (nativeMemoryStoreWord64
                          b"independent-epoch"
                          (nativePointwiseAt stateAt 2)
                          (nativeLaunchRecipeImmediate (nativeTreeEpoch (nitCheckpoint layout)))
                          tail)))
                    (lambda unrestricted previous : Nat .
                      (lambda unrestricted ignored : (family NativePhysicalCommands) .
                        (nativeTreeRestore
                          (nitCheckpoint layout)
                          fanout
                          (niOwners slots system)
                          role
                          (nitOwnOf slots system)
                          (nitPaths (niOwners slots system) role)
                          (lambda unrestricted shard : (family NativeRegionShard) .
                            (nitValuesAt layout))
                          (nativeMemoryCopyOperands
                            (nativeMemoryState stateAt)
                            (nativeMemoryState (nativeTreeRestoredCaptureAt (nitCheckpoint layout)))
                            nativeRegionCaptureBytes
                            nativePeerScratchSlot
                            b"independent-restore-capture"
                            (nativeLaunchRecipeAssertEqual
                              (nativeMemoryLoad stateAt)
                              (nativeMemoryLoad (nativePointwiseAt stateAt 1))
                              b"independent-input-position"
                              (nativeCheckedWordAdd
                                (nativeMemoryState (nativePointwiseAt stateAt 2))
                                (nativeMemoryLoad (nativePointwiseAt stateAt 2))
                                (nativeLaunchRecipeImmediate 1)
                                nativePeerScratchSlot
                                b"independent-epoch-overflow"
                                tail))))))
                    resume)))))))))

-- The publication plan of this generation: the planned layouts with the
-- new membership epoch after a resume.
def nitPublicationPlan =
  (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
    (lambda unrestricted resume : Nat .
      (nativeTreeCheckpointAdvanced (nitCheckpoint layout) resume)))

def nitCommands =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
        (lambda unrestricted fanout : Nat .
          (lambda unrestricted role : Nat .
            (lambda unrestricted edges : (family StdList (family NativePeerEdge)) .
              (lambda unrestricted resume : Nat .
                (lambda unrestricted steps : Nat .
                  (lambda unrestricted step : (pi unrestricted tail : (family NativePhysicalCommands) . (family NativePhysicalCommands)) .
                    (let unrestricted owners =
                      (niOwners slots system)
                      in
                      (let unrestricted stateAt =
                        (nitStateAt layout)
                        in
                        (niInitializeAt
                          slots
                          system
                          (nitValuesAt layout)
                          role
                          (nitStart
                            slots
                            system
                            layout
                            fanout
                            role
                            resume
                            (niFinite
                              (nitValuesAt layout)
                              slots
                              (nativePeerEdgeSetup
                                role
                                edges
                                (nitEndpoint slots layout)
                                (niRoundsAt
                                  stateAt
                                  steps
                                  step
                                  (nativeTreePublish
                                    (nitPublicationPlan layout resume)
                                    fanout
                                    owners
                                    role
                                    edges
                                    (nitCommitmentFor slots system fanout layout)
                                    (nitReleaseFor slots system fanout layout)
                                    stateAt
                                    nativeRegionCaptureBytes
                                    (lambda unrestricted shard : (family NativeRegionShard) .
                                      nativeRegionCaptureBytes)
                                    (nitPaths owners role)
                                    (nativePeerEdgesFold
                                      edges
                                      (lambda unrestricted index : Nat .
                                        (lambda unrestricted edge : (family NativePeerEdge) .
                                        (nativeSocketCloseChecked
                                        (nativeLaunchRecipeSlotValue
                                        (nativePeerEdgeConnectionSlot index))
                                        nativePeerScratchSlot)))
                                      (niOutput stateAt (niCaptureBytes slots system) niHalt)))))))))))))))))))

-- Every edge's frames share the template's storage exactly, and the plan,
-- frames, state and endpoint state fit the declared arena.
def nitFramesAdmitted =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted fanout : Nat .
        (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
          (stdListFold
            Nat
            Nat
            (lambda unrestricted index : Nat .
              (lambda unrestricted rest : Nat .
                (naturalAnd
                  rest
                  (naturalAnd
                    (naturalEqual
                      (frameEnd (nitCommitmentFor slots system fanout layout (succ index)))
                      (frameEnd (nitCommitment layout)))
                    (naturalEqual
                      (frameEnd (nitReleaseFor slots system fanout layout (succ index)))
                      (frameEnd (nitRelease layout)))))))
            1
            (stdListNaturalsBelow (naturalPredecessor (niOwners slots system))))))))

def nitRunAdmitted =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (lambda unrestricted fanout : Nat .
        (lambda unrestricted role : Nat .
          (lambda unrestricted addresses : (family StdList Bytes) .
            (lambda unrestricted resume : Nat .
              (lambda unrestricted epoch : Nat .
                (lambda unrestricted steps : Nat .
                  (naturalAnd
                    (niOwnersAdmitted slots system)
                    (naturalAnd
                      (nativePeerTreeListedAdmitted fanout (niOwners slots system) role addresses)
                      (naturalAnd
                        (naturalLessOrEqual resume 1)
                        (naturalAnd
                          (naturalNonzero epoch)
                          (naturalAnd
                            (naturalBelowTwoPower 64 (naturalAdd epoch resume))
                            (naturalBelowTwoPower 64 steps))))))))))))))

def nativeIndependentTreeArtifact =
  (lambda unrestricted target : (family NativeHostTarget) .
    (lambda unrestricted resources : (family NativeResourceContext) .
      (lambda unrestricted slots : Nat .
        (lambda unrestricted system : (family IndependentStateSystem slots) .
          (lambda unrestricted fanout : Nat .
            (lambda unrestricted role : Nat .
              (lambda unrestricted addresses : (family StdList Bytes) .
                (lambda unrestricted resume : Nat .
                  (lambda unrestricted epoch : Nat .
                    (lambda unrestricted steps : Nat .
                      (lambda unrestricted available : Nat .
                        (let unrestricted edges =
                          (nativePeerTreeEdgesListed fanout (niOwners slots system) role addresses)
                          in
                          (nitWithLayout
                            slots
                            system
                            fanout
                            role
                            epoch
                            0
                            (pointwiseProgramTemporaryValues
                              StdU64
                              slots
                              (independentAssignments slots system))
                            available
                            (lambda unrestricted layout : (family NativeIndependentTreeLayout) .
                              (let unrestricted extent =
                                (nitExtent slots layout edges)
                                in
                                (nat-eliminate
                                  (lambda unrestricted valid : Nat . Bytes)
                                  b""
                                  (lambda unrestricted previous : Nat .
                                    (lambda unrestricted ignored : Bytes .
                                      (eliminate
                                        StdResult
                                        (lambda unrestricted result : (family StdResult Bytes (family NativePhysicalCommands)) .
                                        Bytes)
                                        (planNativePointwiseProgram
                                        slots
                                        (independentAssignments slots system)
                                        (nitValuesAt layout)
                                        (nitEndpoint slots layout)
                                        nativePeerScratchSlot
                                        niEnd)
                                        (branch StdFailure error . b"")
                                        (branch
                                        StdSuccess
                                        step
                                        .
                                        (nativeHostTargetResourceArtifact
                                        target
                                        resources
                                        (niProgramFor
                                        (nativePeerEdgeSlots edges)
                                        extent
                                        (nitCommands
                                        slots
                                        system
                                        layout
                                        fanout
                                        role
                                        edges
                                        resume
                                        steps
                                        (nativePhysicalProgramComposeCommandsAppend step))))))))
                                  (naturalAnd
                                    (nitRunAdmitted
                                      slots
                                      system
                                      fanout
                                      role
                                      addresses
                                      resume
                                      epoch
                                      steps)
                                    (naturalAnd
                                      (nitFramesAdmitted slots system fanout layout)
                                      (naturalAnd
                                        (naturalLessOrEqual extent available)
                                        (arenaCertificate
                                        extent
                                        (constructor
                                        ArenaResidents
                                        ArenaResidentsNext
                                        (constructor
                                        ArenaResident
                                        ArenaResidentValue
                                        b"tree-commitment-frame"
                                        (frameBase (nitCommitment layout))
                                        (naturalSaturatingSubtract
                                        (frameEnd (nitCommitment layout))
                                        (frameBase (nitCommitment layout)))
                                        1)
                                        (constructor
                                        ArenaResidents
                                        ArenaResidentsNext
                                        (constructor
                                        ArenaResident
                                        ArenaResidentValue
                                        b"tree-release-frame"
                                        (frameBase (nitRelease layout))
                                        (naturalSaturatingSubtract
                                        (frameEnd (nitRelease layout))
                                        (frameBase (nitRelease layout)))
                                        1)
                                        (constructor
                                        ArenaResidents
                                        ArenaResidentsNext
                                        (constructor
                                        ArenaResident
                                        ArenaResidentValue
                                        b"tree-owner-state-and-endpoints"
                                        (nitStateAt layout)
                                        (naturalSaturatingSubtract extent (nitStateAt layout))
                                        nativeCheckpointWordBytes)
                                        (nativeTreeResidents (nitCheckpoint layout)))))))))))))))))))))))))