Compile.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
-- | Compilation from 'Imp' programs to BDDs.
module Imp.BDD.Compile
  ( Compiled
  , compile
  ) where

import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Proxy (Proxy(..))
import GHC.TypeLits (KnownSymbol, symbolVal)

import Imp.BDD
import Imp.BDD.Builder
import Imp.BDD.WMC (Weight(..))
import Imp.DSL (Imp(..))

-- | Compiled result: BDD manager, variable weights, Knightian variables and index,
--   and the worlds: each return value with its BDD guard.
type Compiled a = (BDDManager, IntMap Weight, Map String Int, Map a BDD)

-- | State accumulated during compilation.
data CompileState = CompileState
  { csWeights :: !(IntMap Weight)
  , csKnights :: !(Map String (BDD, Int))
  }

-- | Read-only environment for compilation.
newtype CompileEnv = CompileEnv
  { ceTag :: String
  }

initState :: CompileState
initState = CompileState IntMap.empty Map.empty

type CompileM = ReaderT CompileEnv (StateT CompileState BDDM)

-- | Compile an Imp program to its worlds, one guard per return value.
compile :: Ord a => Imp g a -> Compiled a
compile prog =
  let comp = runStateT (runReaderT (compileM prog) (CompileEnv "")) initState
      ((worlds, st), mgr) = runState comp emptyManager
  in (mgr, csWeights st, snd <$> csKnights st, worlds)

-- | Lift a BDD-manager action into the compile monad.
liftBDDM :: BDDM a -> CompileM a
liftBDDM = lift . lift

-- | Resolve a (possibly tag-prefixed) Knightian name.
resolveName :: String -> CompileM String
resolveName baseName = do
  currentTag <- asks ceTag
  return $ if null currentTag then baseName else currentTag ++ "." ++ baseName

-- | Allocate a fresh probabilistic BDD variable with the given Bernoulli weight.
flipVar :: Double -> CompileM BDD
flipVar p = do
  (var, varLabel) <- liftBDDM newVar
  modify' $ \s -> s
    { csWeights = IntMap.insert (unVarLabel varLabel) (Prob p) (csWeights s) }
  return var

-- | Return the BDD variable for a Knightian name, allocating one if it
--   hasn't been seen before.
knightVar :: forall n. KnownSymbol n => CompileM BDD
knightVar = do
  name <- resolveName (symbolVal (Proxy :: Proxy n))
  knights <- gets csKnights
  case Map.lookup name knights of
    Just (var, _) -> return var
    Nothing -> do
      (var, varLabel) <- liftBDDM newVar
      let i = Map.size knights
      modify' $ \s -> s
        { csWeights = IntMap.insert (unVarLabel varLabel) (Knight i) (csWeights s)
        , csKnights = Map.insert name (var, i) (csKnights s)
        }
      return var

-- | Compile a program to worlds.
compileM :: Ord a => Imp g a -> CompileM (Map a BDD)
compileM = \case
  ImpReturn a -> return (Map.singleton a BDDTrue)
  ImpFlip p -> do
    v <- flipVar p
    return (Map.fromList [(True, v), (False, bddNot v)])
  ImpKnight (_ :: Proxy n) -> do
    v <- knightVar @n
    return (Map.fromList [(True, v), (False, bddNot v)])
  ImpObserve b ->
    return (Map.singleton () (if b then BDDTrue else BDDFalse))
  ImpBind m f -> do
    worlds <- compileM m
    conjoined <- mapM (\(a, g) -> traverse (liftBDDM . bddAnd g) =<< compileM (f a))
                      (Map.toList worlds)
    liftBDDM (traverse bddAny (Map.unionsWith (++) [ (: []) <$> c | c <- conjoined ]))
  ImpTag (_ :: Proxy t) inner -> do
    let baseTag = symbolVal (Proxy :: Proxy t)
    resolvedTag <- if null baseTag then asks ceTag else resolveName baseTag
    local (\env -> env { ceTag = resolvedTag }) (compileM inner)
  ImpBranch True  t _ -> compileM t
  ImpBranch False _ f -> compileM f