module ResourcePlanContract

import Platform.Linux.Nvidia.ResourcePlan
import Platform.Linux.Nvidia.PlanHost
import Platform.Linux.Nvidia.LinearStep
import Platform.Linux.Nvidia.ProgramPlan
import Runtime.NativeResourcePlan
import Runtime.NativePhysicalMemory
import Runtime.NativePhysicalProgram
import Runtime.NativeLaunchRecipeLoad
import Runtime.ArenaCertificate
import Execution.ResourceGraph
import Compiler.Planning.Resources
import Std.List
import Std.Foundation
import Std.Natural
import Data.Bytes
import Model.Config
import Model.Parameter
import Model.Word64

-- Hand-computed accounting fixtures. No capacity, page size or alias here
-- is a measured device fact. Two process contexts share one host domain;
-- their independent GPU domains cannot borrow each other's unused capacity.
def rpGraph =
  (lambda unrestricted host : Nat .
    (lambda unrestricted left : Nat .
      (lambda unrestricted right : Nat .
        (constructor
          ResourceGraph
          ResourceGraphOf
          (constructor
            StdList
            StdListCons
            (family ResourceEndpoint)
            (constructor
              ResourceEndpoint
              ResourceEndpointOf
              0
              b"fixture-engine-0"
              b"fixture-native"
              (constructor StdList StdListEmpty Bytes)
              (constructor StdList StdListCons Nat 0 (constructor StdList StdListEmpty Nat)))
            (constructor
              StdList
              StdListCons
              (family ResourceEndpoint)
              (constructor
                ResourceEndpoint
                ResourceEndpointOf
                1
                b"fixture-engine-1"
                b"fixture-native"
                (constructor StdList StdListEmpty Bytes)
                (constructor
                  StdList
                  StdListCons
                  Nat
                  0
                  (constructor StdList StdListCons Nat 1 (constructor StdList StdListEmpty Nat))))
              (constructor
                StdList
                StdListCons
                (family ResourceEndpoint)
                (constructor
                  ResourceEndpoint
                  ResourceEndpointOf
                  2
                  b"fixture-engine-2"
                  b"fixture-native"
                  (constructor StdList StdListEmpty Bytes)
                  (constructor
                    StdList
                    StdListCons
                    Nat
                    0
                    (constructor StdList StdListCons Nat 2 (constructor StdList StdListEmpty Nat))))
                (constructor StdList StdListEmpty (family ResourceEndpoint)))))
          (constructor
            StdList
            StdListCons
            (family ResourceDomain)
            (constructor
              ResourceDomain
              ResourceCapacityDomain
              0
              b"fixture-memory-0"
              (constructor ResourceUnit ResourceMemoryBytes)
              (constructor ResourceFact ResourceEstablished host b"hand-computed-reservation"))
            (constructor
              StdList
              StdListCons
              (family ResourceDomain)
              (constructor
                ResourceDomain
                ResourceCapacityDomain
                1
                b"fixture-memory-1"
                (constructor ResourceUnit ResourceMemoryBytes)
                (constructor ResourceFact ResourceEstablished left b"hand-computed-reservation"))
              (constructor
                StdList
                StdListCons
                (family ResourceDomain)
                (constructor
                  ResourceDomain
                  ResourceCapacityDomain
                  2
                  b"fixture-memory-2"
                  (constructor ResourceUnit ResourceMemoryBytes)
                  (constructor ResourceFact ResourceEstablished right b"hand-computed-reservation"))
                (constructor StdList StdListEmpty (family ResourceDomain)))))
          (constructor StdList StdListEmpty (family ResourceLink))))))

def rpContext =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted gpu : Nat .
      (lambda unrestricted host : Nat .
        (lambda unrestricted left : Nat .
          (lambda unrestricted right : Nat .
            (constructor
              NvidiaResourceContext
              NvidiaResourceContextOf
              (constructor
                NativeResourceContext
                NativeResourceContextOf
                (constructor
                  ResourceSchedule
                  ResourceScheduleOf
                  (rpGraph host left right)
                  (constructor StdList StdListEmpty (family ResourcePlacement))
                  (constructor StdList StdListEmpty (family ResourceLease)))
                identity
                0
                0
                4096
                128
                16384)
              gpu
              gpu
              4096
              4096
              65536))))))

def rpSystem =
  (constructor NvidiaPlanHostMemory NvidiaPlanHostSystemMemory)

def rpVideo =
  (constructor NvidiaPlanHostMemory NvidiaPlanHostVideoMemory)

def rpPaged =
  (constructor NvidiaPlanHostMemory NvidiaPlanHostPagedSystemMemory)

def rpCached =
  (constructor NvidiaPlanHostMemory NvidiaPlanHostGPUCachedSystemMemory)

def rpBuffer =
  (lambda unrestricted identity : Bytes .
    (lambda unrestricted gpu : Nat .
      (lambda unrestricted host : Nat .
        (lambda unrestricted extent : Nat .
          (lambda unrestricted kind : (family NvidiaPlanHostMemory) .
            (constructor
              NvidiaPlanHostBuffer
              NvidiaPlanHostBufferValue
              identity
              gpu
              host
              extent
              kind))))))

-- The fixture's 4 KiB ring buffer holds a 512-entry direct-RM ring (8 bytes
-- an entry); admission refuses a lifecycle whose ring the buffer cannot hold.
def rpDirectRingEntries : Nat =
  512

def rpDirect =
  (lambda unrestricted dataBytes : Nat .
    (lambda unrestricted errorHost : Nat .
      (constructor
        NvidiaPlanHostLayout
        NvidiaPlanHostLayoutValue
        (rpBuffer b"ring" 0x10000 0x60010000 4096 rpSystem)
        (rpBuffer b"push" 0x20000 0x60020000 4096 rpSystem)
        (rpBuffer b"semaphore" 0x30000 0x60030000 4096 rpSystem)
        (rpBuffer b"program" 0x40000 0x60040000 4096 rpSystem)
        (rpBuffer b"qmd" 0x50000 0x60050000 4096 rpSystem)
        (rpBuffer b"userd" 0x60000 0 4096 rpVideo)
        (constructor
          NvidiaPlanHostBuffers
          NvidiaPlanHostBuffersNext
          (rpBuffer b"data" 0x80000 0 dataBytes rpVideo)
          (constructor
            NvidiaPlanHostBuffers
            NvidiaPlanHostBuffersNext
            (rpBuffer b"cached" 0x90000 0 65536 rpCached)
            (constructor
              NvidiaPlanHostBuffers
              NvidiaPlanHostBuffersNext
              (rpBuffer b"paged" 0xa0000 0x600a0000 8192 rpPaged)
              (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersEnd))))
        nvidiaPlanHostAmpere580
        (constructor NvidiaPlanHostLifecycle NvidiaPlanHostDirectRM rpDirectRingEntries)
        b"fixture"
        errorHost)))

-- The UVM channel's GPFIFO buffer: its fixed 0x8000-entry ring and the
-- USERD page after the entries, which the userd entry aliases (offset
-- 0x40000 inside; 0x41000 is past the buffer's end).
def rpUVMRingBytes : Nat =
  (naturalAdd nvidiaPlanHostUVMGPFIFOBytes nvidiaPlanHostPageBytes)

def rpAliased =
  (lambda unrestricted dataBytes : Nat .
    (lambda unrestricted errorHost : Nat .
      (lambda unrestricted aliasOffset : Nat .
        (constructor
          NvidiaPlanHostLayout
          NvidiaPlanHostLayoutValue
          (rpBuffer b"ring" 0x10000 0 rpUVMRingBytes rpVideo)
          (rpBuffer b"push" 0x60000 0x60020000 4096 rpSystem)
          (rpBuffer b"semaphore" 0x70000 0x60030000 4096 rpSystem)
          (rpBuffer b"program" 0x80000 0x60040000 4096 rpSystem)
          (rpBuffer b"qmd" 0x90000 0x60050000 4096 rpSystem)
          (rpBuffer b"userd" (naturalAdd 0x10000 aliasOffset) 0 4096 rpVideo)
          (constructor
            NvidiaPlanHostBuffers
            NvidiaPlanHostBuffersNext
            (rpBuffer b"data" 0xa0000 0 dataBytes rpVideo)
            (constructor
              NvidiaPlanHostBuffers
              NvidiaPlanHostBuffersNext
              (rpBuffer b"cached" 0xb0000 0 65536 rpCached)
              (constructor
                NvidiaPlanHostBuffers
                NvidiaPlanHostBuffersNext
                (rpBuffer b"paged" 0xc0000 0x600a0000 8192 rpPaged)
                (constructor NvidiaPlanHostBuffers NvidiaPlanHostBuffersEnd))))
          nvidiaPlanHostAmpere580
          (constructor NvidiaPlanHostLifecycle NvidiaPlanHostUVM 65536 1048576)
          b"fixture"
          errorHost))))

def rpProgram =
  (nativePhysicalProgramFromCommands
    b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    64
    2
    (nativeLaunchRecipeCommand
      (constructor NativePhysicalOperation NativePhysicalHaltSuccess)
      b"done"
      (constructor NativePhysicalCommands NativePhysicalCommandsEnd)))

def rpLeaseAmount =
  (lambda unrestricted value : (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)) .
    (eliminate
      StdResult
      (lambda unrestricted r : (family StdResult (family NativePhysicalMemoryFailure) (family ResourceLease)) .
        Nat)
      value
      (branch StdFailure error . 0)
      (branch StdSuccess lease . (resourceLeaseAmount lease))))

def rpLeaseTotal =
  (lambda unrestricted domain : Nat .
    (lambda unrestricted leases : (family StdResult (family NativePhysicalMemoryFailure) (family StdList (family ResourceLease))) .
      (eliminate
        StdResult
        (lambda unrestricted r : (family StdResult (family NativePhysicalMemoryFailure) (family StdList (family ResourceLease))) .
          Nat)
        leases
        (branch StdFailure error . 0)
        (branch
          StdSuccess
          values
          .
          (stdListFold
            (family ResourceLease)
            Nat
            (lambda unrestricted lease : (family ResourceLease) .
              (lambda unrestricted total : Nat .
                (naturalAdd
                  (naturalSelect
                    (naturalEqual domain (resourceLeaseDomain lease))
                    (resourceLeaseAmount lease)
                    0)
                  total)))
            0
            values)))))

def rpBoth =
  (lambda unrestricted host : Nat .
    (lambda unrestricted left : Nat .
      (lambda unrestricted right : Nat .
        (lambda unrestricted identity : Bytes .
          (eliminate
            StdResult
            (lambda unrestricted r : (family StdResult (family NativeResourceFailure) (family ResourceSchedule)) .
              Nat)
            (admitNvidiaPlanHostResources
              (rpContext b"first" 1 host left right)
              (rpDirect 8192 0)
              rpProgram)
            (branch StdFailure reason . 0)
            (branch
              StdSuccess
              schedule
              .
              (nvidiaPlanHostResourcesAdmitted
                (nvidiaResourceWithNative
                  (rpContext identity 2 host left right)
                  (nativeResourceContextAddLeases
                    (nvidiaResourceNative (rpContext identity 2 host left right))
                    (resourceScheduleLeases schedule)))
                (rpDirect 8192 0)
                rpProgram)))))))

def rpWithStaging =
  (lambda unrestricted host : Nat .
    (eliminate
      StdResult
      (lambda unrestricted r : (family StdResult (family NativePhysicalMemoryFailure) (family NvidiaResourceContext)) .
        Nat)
      (nlsResourceContext (rpContext b"first" 1 host 12288 1048576))
      (branch StdFailure reason . 0)
      (branch
        StdSuccess
        context
        .
        (nvidiaPlanHostResourcesAdmitted context (rpDirect 8192 0) rpProgram))))

def resultdirectLayout =
  (nvidiaPlanHostLayoutCertified (rpDirect 8192 0) (constructor ArenaResidents ArenaResidentsEnd))

def directLayout : (equal Nat resultdirectLayout 1) =
  (refl Nat 1)

def resultdirectHostCharge =
  (rpLeaseTotal
    0
    (nvidiaPlanHostResourceLeases (rpContext b"first" 1 1048576 1048576 1048576) (rpDirect 8192 0)))

def directHostCharge : (equal Nat resultdirectHostCharge 163840) =
  (refl Nat 163840)

def resultdirectVideoCharge =
  (rpLeaseTotal
    1
    (nvidiaPlanHostResourceLeases (rpContext b"first" 1 1048576 1048576 1048576) (rpDirect 8192 0)))

def directVideoCharge : (equal Nat resultdirectVideoCharge 12288) =
  (refl Nat 12288)

def resultsystemPageCharge =
  (rpLeaseAmount
    (nvidiaResourceBufferLease
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpBuffer b"small" 0xb0000 0 1 rpSystem)))

def systemPageCharge : (equal Nat resultsystemPageCharge 4096) =
  (refl Nat 4096)

def resultpagedSystemCharge =
  (rpLeaseAmount
    (nvidiaResourceBufferLease
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpBuffer b"small" 0xb0000 0 1 rpPaged)))

def pagedSystemCharge : (equal Nat resultpagedSystemCharge 4096) =
  (refl Nat 4096)

def resultcachedSystemCharge =
  (rpLeaseAmount
    (nvidiaResourceBufferLease
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpBuffer b"small" 0xb0000 0 1 rpCached)))

def cachedSystemCharge : (equal Nat resultcachedSystemCharge 65536) =
  (refl Nat 65536)

def resultvideoPageCharge =
  (rpLeaseAmount
    (nvidiaResourceBufferLease
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpBuffer b"small" 0xb0000 0 1 rpVideo)))

def videoPageCharge : (equal Nat resultvideoPageCharge 4096) =
  (refl Nat 4096)

def resultmappedErrorSameCharge =
  (rpLeaseTotal
    0
    (nvidiaPlanHostResourceLeases
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpDirect 8192 0x60150000)))

def mappedErrorSameCharge : (equal Nat resultmappedErrorSameCharge 163840) =
  (refl Nat 163840)

def resulterrorUnmappedStillCharged =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 163840 12288 1048576)
    (rpDirect 8192 0)
    rpProgram)

def errorUnmappedStillCharged : (equal Nat resulterrorUnmappedStillCharged 0) =
  (refl Nat 0)

def resultexactHostAndVideo =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 167936 12288 1048576)
    (rpDirect 8192 0)
    rpProgram)

def exactHostAndVideo : (equal Nat resultexactHostAndVideo 1) =
  (refl Nat 1)

def resultshortHost =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 167935 12288 1048576)
    (rpDirect 8192 0)
    rpProgram)

def shortHost : (equal Nat resultshortHost 0) =
  (refl Nat 0)

def resultshortVideo =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 167936 12287 1048576)
    (rpDirect 8192 0)
    rpProgram)

def shortVideo : (equal Nat resultshortVideo 0) =
  (refl Nat 0)

def resultlargerDeviceStateRefused =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 1048576 12288 1048576)
    (rpDirect 12288 0)
    rpProgram)

def largerDeviceStateRefused : (equal Nat resultlargerDeviceStateRefused 0) =
  (refl Nat 0)

def resultuvmLayout =
  (nvidiaPlanHostLayoutCertified
    (rpAliased 8192 0 0x40000)
    (constructor ArenaResidents ArenaResidentsEnd))

def uvmLayout : (equal Nat resultuvmLayout 1) =
  (refl Nat 1)

def resultuvmHostCharge =
  (rpLeaseTotal
    0
    (nvidiaPlanHostResourceLeases
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpAliased 8192 0 0x40000)))

def uvmHostCharge : (equal Nat resultuvmHostCharge 159744) =
  (refl Nat 159744)

def resultuvmVideoAliasChargedOnce =
  (rpLeaseTotal
    1
    (nvidiaPlanHostResourceLeases
      (rpContext b"first" 1 1048576 1048576 1048576)
      (rpAliased 8192 0 0x40000)))

def uvmVideoAliasChargedOnce : (equal Nat resultuvmVideoAliasChargedOnce 274432) =
  (refl Nat 274432)

def resultuvmExactReservation =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 163840 274432 1048576)
    (rpAliased 8192 0 0x40000)
    rpProgram)

def uvmExactReservation : (equal Nat resultuvmExactReservation 1) =
  (refl Nat 1)

-- one video byte fewer than the ring, its USERD page and the data is refused
def resultuvmShortVideo =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 163840 274431 1048576)
    (rpAliased 8192 0 0x40000)
    rpProgram)

def uvmShortVideo : (equal Nat resultuvmShortVideo 0) =
  (refl Nat 0)

def resultuvmAliasOutsideRefused =
  (nvidiaPlanHostResourcesAdmitted
    (rpContext b"first" 1 1048576 1048576 1048576)
    (rpAliased 8192 0 0x41000)
    rpProgram)

def uvmAliasOutsideRefused : (equal Nat resultuvmAliasOutsideRefused 0) =
  (refl Nat 0)

def resulttwoDevicesExact =
  (rpBoth 335872 12288 12288 b"second")

def twoDevicesExact : (equal Nat resulttwoDevicesExact 1) =
  (refl Nat 1)

def resultsharedHostOneByteShort =
  (rpBoth 335871 12288 12288 b"second")

def sharedHostOneByteShort : (equal Nat resultsharedHostOneByteShort 0) =
  (refl Nat 0)

def resultfirstCannotBorrowSecondVideo =
  (rpBoth 1048576 4096 28672 b"second")

def firstCannotBorrowSecondVideo : (equal Nat resultfirstCannotBorrowSecondVideo 0) =
  (refl Nat 0)

def resultsecondCannotBorrowFirstVideo =
  (rpBoth 1048576 28672 4096 b"second")

def secondCannotBorrowFirstVideo : (equal Nat resultsecondCannotBorrowFirstVideo 0) =
  (refl Nat 0)

def resultduplicateAllocationNamesRefused =
  (rpBoth 1048576 1048576 1048576 b"first")

def duplicateAllocationNamesRefused : (equal Nat resultduplicateAllocationNamesRefused 0) =
  (refl Nat 0)

def resultzeroMappingRefused =
  (stdBoolToNatural
    (stdResultIsSuccess (family NativePhysicalMemoryFailure) Nat (nativeMemoryMappedBytes 4096 0)))

def zeroMappingRefused : (equal Nat resultzeroMappingRefused 0) =
  (refl Nat 0)

def resultwordRoundingOverflowRefused =
  (stdBoolToNatural
    (stdResultIsSuccess
      (family NativePhysicalMemoryFailure)
      Nat
      (nativeMemoryMappedBytes 4096 18446744073709551615)))

def wordRoundingOverflowRefused : (equal Nat resultwordRoundingOverflowRefused 0) =
  (refl Nat 0)

def resultstagingIncluded =
  (rpWithStaging 16945152)

def stagingIncluded : (equal Nat resultstagingIncluded 1) =
  (refl Nat 1)

def resultstagingOneByteShort =
  (rpWithStaging 16945151)

def stagingOneByteShort : (equal Nat resultstagingOneByteShort 0) =
  (refl Nat 0)

def observations =
  (bytes-append
    (dataBytesWord64LE (modelWord64FromNaturalTruncated resultdirectLayout))
    (bytes-append
      (dataBytesWord64LE (modelWord64FromNaturalTruncated resultdirectHostCharge))
      (bytes-append
        (dataBytesWord64LE (modelWord64FromNaturalTruncated resultdirectVideoCharge))
        (bytes-append
          (dataBytesWord64LE (modelWord64FromNaturalTruncated resultsystemPageCharge))
          (bytes-append
            (dataBytesWord64LE (modelWord64FromNaturalTruncated resultpagedSystemCharge))
            (bytes-append
              (dataBytesWord64LE (modelWord64FromNaturalTruncated resultcachedSystemCharge))
              (bytes-append
                (dataBytesWord64LE (modelWord64FromNaturalTruncated resultvideoPageCharge))
                (bytes-append
                  (dataBytesWord64LE (modelWord64FromNaturalTruncated resultmappedErrorSameCharge))
                  (bytes-append
                    (dataBytesWord64LE
                      (modelWord64FromNaturalTruncated resulterrorUnmappedStillCharged))
                    (bytes-append
                      (dataBytesWord64LE (modelWord64FromNaturalTruncated resultexactHostAndVideo))
                      (bytes-append
                        (dataBytesWord64LE (modelWord64FromNaturalTruncated resultshortHost))
                        (bytes-append
                          (dataBytesWord64LE (modelWord64FromNaturalTruncated resultshortVideo))
                          (bytes-append
                            (dataBytesWord64LE
                              (modelWord64FromNaturalTruncated resultlargerDeviceStateRefused))
                            (bytes-append
                              (dataBytesWord64LE (modelWord64FromNaturalTruncated resultuvmLayout))
                              (bytes-append
                                (dataBytesWord64LE
                                  (modelWord64FromNaturalTruncated resultuvmHostCharge))
                                (bytes-append
                                  (dataBytesWord64LE
                                    (modelWord64FromNaturalTruncated resultuvmVideoAliasChargedOnce))
                                  (bytes-append
                                    (dataBytesWord64LE
                                      (modelWord64FromNaturalTruncated resultuvmExactReservation))
                                    (bytes-append
                                      (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultuvmAliasOutsideRefused))
                                      (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated resulttwoDevicesExact))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultsharedHostOneByteShort))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultfirstCannotBorrowSecondVideo))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultsecondCannotBorrowFirstVideo))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultduplicateAllocationNamesRefused))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated resultzeroMappingRefused))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated
                                        resultwordRoundingOverflowRefused))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated resultstagingIncluded))
                                        (bytes-append
                                        (dataBytesWord64LE
                                        (modelWord64FromNaturalTruncated resultstagingOneByteShort))
                                        b"")))))))))))))))))))))))))))

def main : Nat =
  0