TemporaryFunctions.hs

{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RebindableSyntax    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Tests.Control.TemporaryFunctions where

-- base
import Data.Maybe
  ( fromJust )

-- vector-sized
import qualified Data.Vector.Sized as Vector

-- fir
import FIR

------------------------------------------------
-- program

-- Indexing a constant array with a run-time index requires a local variable
-- holding the array. Each function must use its own local variable.

table :: Array 4 Float
table = MkArray ( fromJust $ Vector.fromList [ 1, 2, 3, 4 ] )

type Defs =
  '[ "ubo"  ':-> Uniform '[ DescriptorSet 0, Binding 0 ]
                   ( Struct '[ "i" ':-> Word32 ] )
   , "out"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                   ( Struct '[ "res" ':-> Float ] )
   , "f"    ':-> Function '[ "u" ':-> Var R Word32 ] Float
   , "g"    ':-> Function '[ "v" ':-> Var R Word32 ] Float
   , "main" ':-> EntryPoint '[ LocalSize 1 1 1 ] Compute
   ]

program :: Module Defs
program = Module do

  ( f :: Code Word32 -> Code Float ) <- fundef @"f" do
    u <- get @"u"
    pure ( view @( AnIndex ( Code Word32 ) ) u ( Lit table :: Code ( Array 4 Float ) ) )

  ( g :: Code Word32 -> Code Float ) <- fundef @"g" do
    v <- get @"v"
    pure ( view @( AnIndex ( Code Word32 ) ) ( v + 1 ) ( Lit table :: Code ( Array 4 Float ) ) )

  entryPoint @"main" @Compute do
    i <- use @( Name "ubo" :.: Name "i" )
    assign @( Name "out" :.: Name "res" ) ( f i + g i )