Use Set instead of Map for underlying type of Sigma

A minor shortcoming is that indeed the resulting rows are ordered in a less
pleasing way: previously, each row can be interpreted as a little-endian
number, and the rows are in descending order.

Author: kccqzy

package.yaml

@@ -11,7 +11,7 @@ dependencies:
 - mtl
 - transformers
 - errors
-- containers
+- containers >= 0.5.11
 - array
 - reflection
 - megaparsec

src/Prob/CoreAST.hs

@@ -9,9 +9,10 @@ module Prob.CoreAST
   , Stmt(..)
   , Prog(..)
   , Sigma
+  , sigmaInsert
   ) where
 
-import qualified Data.Map.Strict as M
+import qualified Data.Set as Set
 import Data.String
 
 -- | The syntax of expressions, parametrized by the representation of variables
@@ -59,5 +60,10 @@ deriving instance Foldable (Prog r)
 instance IsString (Expr String) where
   fromString = Var
 
--- | Sigma is just the set of all variables assignments.
-type Sigma vt = M.Map vt Bool
+-- | Sigma is just the set of all variables assignments. Since our language only
+-- ever deals with Bool variables, we use a 'Set.Set' and the presence/absence
+-- indicates their values.
+type Sigma vt = Set.Set vt
+
+sigmaInsert :: Ord vt => vt -> Bool -> Sigma vt -> Sigma vt
+sigmaInsert x v = (if v then Set.insert else Set.delete) x

src/Prob/Den.hs

@@ -15,7 +15,6 @@ module Prob.Den
   , runDenStmt
   ) where
 
-import Control.Error
 import Control.Monad.Reader
 import Control.Monad.State
 import Data.Bifunctor
@@ -30,11 +29,11 @@ import qualified Prob.LinearEq as L
 -- Denotational Semantics
 --------------------------------------------------------------------------------
 
-allPossibleStates :: (Ord k, Foldable t) => t k -> [M.Map k Bool]
-allPossibleStates = foldr (\var -> concatMap (\st -> [M.insert var True st, M.insert var False st])) [M.empty]
+allPossibleStates :: (Ord k) => Set.Set k -> [Sigma k]
+allPossibleStates = Set.toList . Set.powerSet
 
 denExpr :: (Show vt, Ord vt) => Expr vt -> Sigma vt -> Bool
-denExpr (Var x) sigma = fromMaybe (error $ "undefined variable " ++ show x) $ M.lookup x sigma
+denExpr (Var x) sigma = Set.member x sigma
 denExpr (Constant d) _ = d
 denExpr (Or a b) sigma = denExpr a sigma || denExpr b sigma
 denExpr (And a b) sigma = denExpr a sigma && denExpr b sigma
@@ -52,10 +51,10 @@ type Den vt = StateT (Maybe (CurrentLoop vt)) (Reader (Sigma vt))
 
 denStmt :: (Show vt, Ord vt) => [Stmt vt] -> Sigma vt -> Den vt (L.RHS (Sigma vt))
 denStmt [] sigma = lift $ ReaderT $ \sigma' -> if sigma' == sigma then pure (L.RHS 1 []) else pure (L.RHS 0 [])
-denStmt ((x := e):next) sigma = denStmt next (M.insert x (denExpr e sigma) sigma)
+denStmt ((x := e):next) sigma = denStmt next (sigmaInsert x (denExpr e sigma) sigma)
 denStmt ((x :~ Bernoulli theta):next) sigma = do
-  dTrue <- denStmt next (M.insert x True sigma)
-  dFalse <- denStmt next (M.insert x False sigma)
+  dTrue <- denStmt next (sigmaInsert x True sigma)
+  dFalse <- denStmt next (sigmaInsert x False sigma)
   pure ((theta `mult` dTrue) `plus` ((1 - theta) `mult` dFalse))
   where mult :: Rational -> L.RHS x -> L.RHS x
         mult k (L.RHS c tms) = L.RHS (k * c) (map (\(L.Term b y) -> L.Term (k*b) y) tms)
@@ -98,7 +97,7 @@ findDenProg :: (Ord vt) => [vt] -> (Set.Set vt -> Sigma vt -> r) -> r
 findDenProg p g = g vars initialState
   where
     vars = Set.fromList p
-    initialState = M.fromSet (const False) vars
+    initialState = Set.empty
                    -- initial state: all variables initialized to False
 
 extractRHS :: L.RHS vt -> Rational

src/Prob/Eval.hs

@@ -15,8 +15,8 @@ import Control.Monad.ST
 import Control.Monad.State
 import Data.Bifunctor
 import Data.List
-import qualified Data.Map.Strict as M
 import Data.Ratio
+import qualified Data.Set as Set
 import Prob.CoreAST
 import System.Random.MWC
 import System.Random.MWC.Distributions
@@ -32,13 +32,13 @@ type ProgState vt s = (Sigma vt, Gen s)
 type Eval vt s = MaybeT (StateT (ProgState vt s) (ST s))
 
 runE :: Eval vt s a -> IO (Maybe a)
-runE e = withSystemRandom . asGenST $ (\rng -> evalStateT (runMaybeT e) (M.empty, rng))
+runE e = withSystemRandom . asGenST $ (\rng -> evalStateT (runMaybeT e) (Set.empty, rng))
 
 runEs :: Int -> Eval vt s a -> IO [a]
-runEs t e = withSystemRandom . asGenST $ (\rng -> catMaybes <$> evalStateT (replicateM t (runMaybeT e)) (M.empty, rng))
+runEs t e = withSystemRandom . asGenST $ (\rng -> catMaybes <$> evalStateT (replicateM t (runMaybeT e)) (Set.empty, rng))
 
 evalExpr :: (Show vt, Ord vt) => Expr vt -> Eval vt s Bool
-evalExpr (Var x) = fromMaybe (error $ "undefined variable " ++ show x) <$> gets (M.lookup x . fst)
+evalExpr (Var x) = gets (Set.member x . fst)
 evalExpr (Constant d) = pure d
 evalExpr (Or a b) = liftA2 (||) (evalExpr a) (evalExpr b)
 evalExpr (And a b) = liftA2 (&&) (evalExpr a) (evalExpr b)
@@ -54,11 +54,11 @@ evalStmt :: (Show vt, Ord vt) => [Stmt vt] -> Eval vt s ()
 evalStmt [] = pure ()
 evalStmt ((x := a):next) = do
   v <- evalExpr a
-  modify (first (M.insert x v))
+  modify (first (sigmaInsert x v))
   evalStmt next
 evalStmt ((x :~ d):next) = do
   v <- drawDist d
-  modify (first (M.insert x v))
+  modify (first (sigmaInsert x v))
   evalStmt next
 evalStmt (Observe e:next) = do
   e' <- evalExpr e

src/Prob/Pretty.hs

@@ -7,8 +7,9 @@ module Prob.Pretty
   ) where
 
 import Data.Bifunctor
-import qualified Data.Map.Strict as M
+import Data.Foldable
 import Data.Ratio
+import qualified Data.Set as Set
 import Prob.CoreAST
 import Prob.Den (denProg)
 import Prob.Eval (sampled)
@@ -18,24 +19,28 @@ data Mode = ModeDen | ModeEval Int
 handleProgPretty :: forall vt r. (Show vt, Ord vt) => Prog r vt -> Mode -> IO ShowS
 handleProgPretty p m = formatResult <$> r
   where
+    allVars :: Set.Set vt
+    allVars =
+      Set.fromList $ case p of Return s e -> concatMap toList s ++ toList e; ReturnAll s -> concatMap toList s
     r :: IO [(String, String)]
     r =
       case p of
         Return {} -> map (bimap (`shows` " ") (($ []) . formatRational)) <$> (case m of ModeDen -> pure (denProg p); ModeEval t -> sampled t p)
         ReturnAll {} -> map (bimap pprMap (($ []) . formatRational)) <$> (case m of ModeDen -> pure (denProg p); ModeEval t -> sampled t p)
       where
         pprMap :: Sigma vt -> String
-        pprMap =
-          M.foldrWithKey
-            (\var val s ->
+        pprMap sigma =
+          foldr
+            (\var s ->
                shows var .
                showString " ->" .
                showString
-                 (if val
+                 (if Set.member var sigma
                     then "  true "
                     else " false ") $
                s)
             " "
+            allVars
         formatRational rat = shows (numerator rat) . showChar '/' . shows (denominator rat)
     formatResult :: [(String, String)] -> ShowS
     formatResult [] = showString "No results produced.\n"

stack.yaml

@@ -1,4 +1,4 @@
-resolver: lts-11.8
+resolver: lts-12.26
 packages:
 - .
 extra-deps: []