module Platform.Linux.Nvidia.ResourcePlan

import Platform.Linux.Nvidia.PlanHost
import Runtime.NativePhysicalMemory
import Runtime.NativePhysicalProgram
import Runtime.NativePhysicalNativeELF
import Runtime.NativeResourcePlan
import Runtime.ArenaCertificate
import Execution.ResourceGraph
import Std.Foundation
import Std.List
import Std.Natural
import Data.Bytes

-- Quanta are explicit deployment accounting contracts, not estimates inferred
-- from another card. System, video and GPU-cached system allocations may have
-- different physical granularity. Their provenance remains target/deployment
-- responsibility. GPU virtual addresses do not determine the backing domain.
family NvidiaResourceContext : Type 0
constructor NvidiaResourceContextOf
field unrestricted nvidiaResourceContextNative : (family NativeResourceContext)
field unrestricted nvidiaResourceContextEndpoint : Nat
field unrestricted nvidiaResourceContextVideoDomain : Nat
field unrestricted nvidiaResourceContextSystemQuantum : Nat
field unrestricted nvidiaResourceContextVideoQuantum : Nat
field unrestricted nvidiaResourceContextCachedQuantum : Nat

end-family

def nvidiaResourceNative =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (eliminate
      NvidiaResourceContext
      (lambda unrestricted current : (family NvidiaResourceContext) .
        (family NativeResourceContext))
      context
      (branch
        NvidiaResourceContextOf
        native
        endpoint
        video
        systemQuantum
        videoQuantum
        cachedQuantum
        .
        native)))

def nvidiaResourceWithNative =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted replacement : (family NativeResourceContext) .
      (eliminate
        NvidiaResourceContext
        (lambda unrestricted current : (family NvidiaResourceContext) .
          (family NvidiaResourceContext))
        context
        (branch
          NvidiaResourceContextOf
          native
          endpoint
          video
          systemQuantum
          videoQuantum
          cachedQuantum
          .
          (constructor
            NvidiaResourceContext
            NvidiaResourceContextOf
            replacement
            endpoint
            video
            systemQuantum
            videoQuantum
            cachedQuantum)))))

-- Host-only anonymous mappings use the host page quantum. RM buffers use the
-- target's allocation quanta, even if no CPU mapping of the buffer exists.
def nvidiaResourceHostLease =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted suffix : Bytes .
      (lambda unrestricted extent : Nat .
        (eliminate
          NativeResourceContext
          (lambda unrestricted current : (family NativeResourceContext) .
            (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)))
          (nvidiaResourceNative context)
          (branch
            NativeResourceContextOf
            schedule
            identity
            endpoint
            domain
            quantum
            entries
            pairs
            .
            (nativeMemoryLease (bytes-append identity suffix) endpoint domain quantum extent))))))

def nvidiaResourceBufferLease =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted buffer : (family NvidiaPlanHostBuffer) .
      (eliminate
        NvidiaResourceContext
        (lambda unrestricted current : (family NvidiaResourceContext) .
          (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)))
        context
        (branch
          NvidiaResourceContextOf
          native
          gpu
          video
          systemQuantum
          videoQuantum
          cachedQuantum
          .
          (eliminate
            NativeResourceContext
            (lambda unrestricted current : (family NativeResourceContext) .
              (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)))
            native
            (branch
              NativeResourceContextOf
              schedule
              identity
              host
              system
              quantum
              entries
              pairs
              .
              (nativeMemoryLease
                (bytes-append identity (bytes-append b"/buffer/" (phBufferIdentity buffer)))
                gpu
                (naturalSelect (phBufferIsVideo buffer) video system)
                (naturalSelect
                  (phBufferIsVideo buffer)
                  videoQuantum
                  (naturalSelect (phBufferIsGPUCached buffer) cachedQuantum systemQuantum))
                (phBufferExtent buffer))))))))

def nvidiaResourcePrependLease =
  (lambda unrestricted head : (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)) .
    (lambda unrestricted rest : (family StdResult (family NativePhysicalMemoryFailure) (family StdList (family ResourceLease))) .
      (stdResultBind
        (family NativePhysicalMemoryFailure)
        (family ResourceLease)
        (family StdList (family ResourceLease))
        head
        (lambda unrestricted lease : (family ResourceLease) .
          (stdResultBind
            (family NativePhysicalMemoryFailure)
            (family StdList (family ResourceLease))
            (family StdList (family ResourceLease))
            rest
            (lambda unrestricted leases : (family StdList (family ResourceLease)) .
              (constructor
                StdResult
                StdSuccess
                (family NativePhysicalMemoryFailure)
                (family StdList (family ResourceLease))
                (constructor StdList StdListCons (family ResourceLease) lease leases))))))))

-- phAllocatedBuffers owns the direct-RM/UVM allocation distinction. In UVM
-- the USERD range aliases GPFIFO and contributes no second allocation. The
-- error notifier is always allocated, even when errorHost is zero. It has no
-- GPU VA and is therefore not one of phAllocatedBuffers' placed objects.
-- A CPU view of a device buffer and the MMIO doorbell consume no second RAM
-- allocation. Page-table/driver overhead is outside this explicit inventory.
def nvidiaPlanHostResourceLeases =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
      (nvidiaResourcePrependLease
        (nvidiaResourceHostLease context b"/control/parameters" nvidiaPlanHostParameterMappingBytes)
        (eliminate
          NvidiaPlanHostBuffers
          (lambda unrestricted current : (family NvidiaPlanHostBuffers) .
            (family
              StdResult
              (family NativePhysicalMemoryFailure)
              (family StdList (family ResourceLease))))
          (phAllocatedBuffers layout)
          (branch
            NvidiaPlanHostBuffersEnd
            .
            (nvidiaResourcePrependLease
              (nvidiaResourceBufferLease
                context
                (constructor
                  NvidiaPlanHostBuffer
                  NvidiaPlanHostBufferValue
                  b"error-notifier"
                  0
                  0
                  nvidiaPlanHostErrorPageBytes
                  (constructor NvidiaPlanHostMemory NvidiaPlanHostSystemMemory)))
              (constructor
                StdResult
                StdSuccess
                (family NativePhysicalMemoryFailure)
                (family StdList (family ResourceLease))
                (constructor StdList StdListEmpty (family ResourceLease)))))
          (branch
            NvidiaPlanHostBuffersNext
            head
            rest
            induction
            .
            (nvidiaResourcePrependLease (nvidiaResourceBufferLease context head) induction))))))

def admitNvidiaPlanHostResources =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
      (lambda unrestricted program : (family NativePhysicalProgram) .
        (eliminate
          StdBool
          (lambda unrestricted current : (family StdBool) .
            (family StdResult (family NativeResourceFailure) (family ResourceSchedule)))
          (stdBoolFromNatural
            (naturalAnd
              (nvidiaPlanHostLayoutAdmitted layout)
              (nvidiaPlanHostLayoutCertified layout (constructor ArenaResidents ArenaResidentsEnd))))
          (branch
            StdTrue
            .
            (eliminate
              StdResult
              (lambda unrestricted current : (family StdResult (family NativePhysicalMemoryFailure) (family StdList (family ResourceLease))) .
                (family StdResult (family NativeResourceFailure) (family ResourceSchedule)))
              (nvidiaPlanHostResourceLeases context layout)
              (branch
                StdFailure
                reason
                .
                (constructor
                  StdResult
                  StdFailure
                  (family NativeResourceFailure)
                  (family ResourceSchedule)
                  (constructor NativeResourceFailure NativeResourceMemoryFailed reason)))
              (branch
                StdSuccess
                leases
                .
                (admitNativePhysicalResources
                  (nativeResourceContextAddLeases (nvidiaResourceNative context) leases)
                  program))))
          (branch
            StdFalse
            .
            (constructor
              StdResult
              StdFailure
              (family NativeResourceFailure)
              (family ResourceSchedule)
              (constructor
                NativeResourceFailure
                NativeResourceAdmissionFailed
                (constructor ResourceFailure ResourceInvalidLease))))))))

def nvidiaPlanHostResourcesAdmitted =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
      (lambda unrestricted program : (family NativePhysicalProgram) .
        (stdBoolToNatural
          (stdResultIsSuccess
            (family NativeResourceFailure)
            (family ResourceSchedule)
            (admitNvidiaPlanHostResources context layout program))))))

def nvidiaPlanHostGenerateResourceELF =
  (lambda unrestricted context : (family NvidiaResourceContext) .
    (lambda unrestricted layout : (family NvidiaPlanHostLayout) .
      (lambda unrestricted program : (family NativePhysicalProgram) .
        (nat-eliminate
          (lambda unrestricted admitted : Nat . Bytes)
          b""
          (lambda unrestricted prior : Nat .
            (lambda unrestricted ignored : Bytes .
              (nativePhysicalGenerateNativeELFBytesDirect program)))
          (nvidiaPlanHostResourcesAdmitted context layout program)))))