Phi.hs

{-# LANGUAGE ScopedTypeVariables #-}

{-|
Module: CodeGen.Phi

Code generation for ϕ-functions, which enables definition of variables in branches when using an SSA form.

See also "CodeGen.CFG" which does code generation for the control flow operation themselves.
-}

module CodeGen.Phi
  ( phiInstruction, phiInstructions )
  where

-- base
import Data.List
  ( foldl1' )
import Data.List.NonEmpty
  ( NonEmpty((:|)) )
import qualified Data.List.NonEmpty as NE
  ( head, singleton, toList )

-- containers
import Data.Map
  ( Map )
import qualified Data.Map.Strict       as Map
import qualified Data.Map.Merge.Strict as Map

-- text-short
import Data.Text.Short
  ( ShortText )


-- fir
import CodeGen.Binary
  ( instruction )
import CodeGen.IDs
  ( typeID )
import CodeGen.Instruction
  ( ID, Pair(Pair), Instruction(..), toArgs )
import CodeGen.Monad
  ( CGMonad, MonadFresh(fresh) )
import qualified SPIRV.Operation as SPIRV.Op
import qualified SPIRV.PrimTy    as SPIRV

----------------------------------------------------------------------------
-- ϕ-instructions

conflicts :: forall k a. (Ord k, Eq a)
          => ( k -> Bool )
          -> [ Map k a ]
          -> Map k (NonEmpty a)
conflicts keyIsOK
  = Map.filterWithKey ( \k as -> keyIsOK k && any (/= NE.head as) as )
  . foldl1'
      ( Map.merge
          Map.dropMissing
          Map.dropMissing
          ( Map.zipWithMatched (const (<>)) )
      )
  . map (fmap NE.singleton)

-- use "Pair" for the correct "Put" instance for list of pairs
phiInstruction :: (ID, SPIRV.PrimTy) -> [ Pair ID ID ] -> CGMonad ()
phiInstruction (v, ty) bdAndBlockIDs
  = do
      tyID <- typeID ty
      instruction
        Instruction
          { operation = SPIRV.Op.Phi
          , resTy     = Just tyID
          , resID     = Just v
          , args      = toArgs bdAndBlockIDs
          }

phiInstructions :: ( ShortText -> Bool )
                -> [ ID ]
                -> [ Map ShortText (ID, SPIRV.PrimTy) ]
                -> CGMonad (Map ShortText ID)
phiInstructions isRelevant blocks bindings 
  = Map.traverseWithKey
      ( \ _ idsAndTys@( (_,ty) :| _ ) -> do
        let bdAndBlockIDs :: [Pair ID ID]
            bdAndBlockIDs
              = zipWith
                  (\(x_ID, _) blk -> Pair (x_ID, blk))
                  (NE.toList idsAndTys)
                  blocks
        v <- fresh
        phiInstruction (v,ty) bdAndBlockIDs
        pure v
      )
      ( conflicts isRelevant bindings )