Selection.hs

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

module Tests.Layout.Selection where

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

-- fir
import FIR

------------------------------------------------
-- selecting between a value loaded from a buffer
-- and a value constructed in the shader

type S = Struct '[ "a" ':-> Float, "b" ':-> Array 2 Float ]

type Defs =
  '[ "ubo"  ':-> Uniform '[ DescriptorSet 0, Binding 0 ]
                   ( Struct '[ "s" ':-> S, "x" ':-> Float ] )
   , "out"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                   ( Struct '[ "s" ':-> S ] )
   , "main" ':-> EntryPoint '[ LocalSize 1 1 1 ] Compute
   ]

program :: Module Defs
program = Module do
  entryPoint @"main" @Compute do
    loaded <- use @( Name "ubo" :.: Name "s" )
    x      <- use @( Name "ubo" :.: Name "x" )
    let built = Struct ( x :& Lit ( MkArray ( Vector.replicate 3 ) ) :& End ) :: Code S
    assign @( Name "out" :.: Name "s" ) ( if x > 0 then loaded else built )