diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2016, The Hakaru Team +Copyright (c) 2017, The Hakaru Team All rights reserved. diff --git a/commands/Compile.hs b/commands/Compile.hs --- a/commands/Compile.hs +++ b/commands/Compile.hs @@ -1,10 +1,12 @@ -{-# LANGUAGE OverloadedStrings, +{-# LANGUAGE CPP, + OverloadedStrings, PatternGuards, DataKinds, KindSignatures, GADTs, FlexibleContexts, - TypeOperators + TypeOperators, + RankNTypes #-} module Main where @@ -19,31 +21,69 @@ import Language.Hakaru.Types.DataKind import Language.Hakaru.Pretty.Haskell -import Language.Hakaru.Command (parseAndInfer) +import Language.Hakaru.Command -import Data.Text -import qualified Data.Text.IO as IO -import System.IO (stderr) -import Text.PrettyPrint +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif -import System.Environment +import Data.Text as TxT +import qualified Data.Text.IO as IO +import Data.Maybe (fromJust) +import Data.Monoid ((<>)) +import System.IO (stderr) +import Text.PrettyPrint hiding ((<>)) +import qualified Options.Applicative as O import System.FilePath -data Options = Options - { program :: String - , asFunc :: Maybe String - } +data Options = + Options { fileIn :: String + , fileOut :: Maybe String + , asModule :: Maybe String + , fileIn2 :: Maybe String + , logFloatPrelude :: Bool + , optimize :: Bool + } deriving Show + main :: IO () main = do - args <- getArgs - case args of - [prog1, prog2] -> compileRandomWalk prog1 prog2 - [prog] -> compileHakaru prog - _ -> IO.hPutStrLn stderr - "Usage: compile \n\ - \ | compile " + opts <- parseOpts + case fileIn2 opts of + Just _ -> compileRandomWalk opts + Nothing -> compileHakaru opts +parseOpts :: IO Options +parseOpts = O.execParser $ O.info (O.helper <*> options) + $ O.fullDesc <> O.progDesc "Compile Hakaru to Haskell" + +{- + +There is some redundancy in Compile.hs, Hakaru.hs, and HKC.hs in how we decide +what to run given which arguments. There may be a way to unify these. + +-} + +options :: O.Parser Options +options = Options + <$> O.strArgument (O.metavar "INPUT" <> + O.help "Program to be compiled" ) + <*> (O.optional $ + O.strOption (O.metavar "FILE" <> + O.short 'o' <> + O.help "Optional output file name")) + <*> (O.optional $ + O.strOption (O.long "as-module" <> + O.short 'M' <> + O.help "creates a haskell module with this name")) + <*> (O.optional $ + O.strOption (O.long "transition-kernel" <> + O.help "Use this program as transition kernel for running a markov chain")) + <*> O.switch (O.long "logfloat-prelude" <> + O.help "use logfloat prelude for numeric stability") + <*> O.switch (O.short 'O' <> + O.help "perform Hakaru AST optimizations") + prettyProg :: (ABT T.Term abt) => String -> abt '[] a @@ -54,62 +94,100 @@ , nest 2 (pretty ast) ]) -compileHakaru :: String -> IO () -compileHakaru file = do - prog <- IO.readFile file - let target = replaceFileName file (takeBaseName file) ++ ".hs" +compileHakaru + :: Options + -> IO () +compileHakaru opts = do + let file = fileIn opts + prog <- readFromFile file case parseAndInfer prog of Left err -> IO.hPutStrLn stderr err Right (TypedAST typ ast) -> do - IO.writeFile target . Data.Text.unlines $ - header ++ - [ pack $ prettyProg "prog" (et ast) ] ++ - footer typ + let ast' = (if optimize opts then optimizations else id) (et ast) + writeHkHsToFile file (fileOut opts) . TxT.unlines $ + header (logFloatPrelude opts) (asModule opts) ++ + [ pack $ prettyProg "prog" ast' ] ++ + (case asModule opts of + Nothing -> footer (logFloatPrelude opts) typ + Just _ -> []) where et = expandTransformations -compileRandomWalk :: String -> String -> IO () -compileRandomWalk f1 f2 = do - p1 <- IO.readFile f1 - p2 <- IO.readFile f2 - let output = replaceFileName f1 (takeBaseName f1) ++ ".hs" +compileRandomWalk + :: Options + -> IO () +compileRandomWalk opts = do + let f1 = fileIn opts + f2 = fromJust . fileIn2 $ opts + p1 <- readFromFile f1 + p2 <- readFromFile f2 case (parseAndInfer p1, parseAndInfer p2) of (Right (TypedAST typ1 ast1), Right (TypedAST typ2 ast2)) -> -- TODO: Use better error messages for type mismatch case (typ1, typ2) of (SFun a (SMeasure b), SMeasure c) | (Just Refl, Just Refl) <- (jmEq1 a b, jmEq1 b c) - -> IO.writeFile output . Data.Text.unlines $ - header ++ + -> writeHkHsToFile f1 (fileOut opts) . TxT.unlines $ + header (logFloatPrelude opts) (asModule opts) ++ [ pack $ prettyProg "prog1" (expandTransformations ast1) ] ++ [ pack $ prettyProg "prog2" (expandTransformations ast2) ] ++ - footerWalk + (case asModule opts of + Nothing -> footerWalk + Just _ -> []) _ -> IO.hPutStrLn stderr "compile: programs have wrong type" (Left err, _) -> IO.hPutStrLn stderr err (_, Left err) -> IO.hPutStrLn stderr err -header :: [Text] -header = +writeHkHsToFile :: String -> Maybe String -> Text -> IO () +writeHkHsToFile inFile moutFile content = + let outFile = case moutFile of + Nothing -> replaceFileName inFile (takeBaseName inFile) ++ ".hs" + Just x -> x + in writeToFile outFile content + +header :: Bool -> Maybe String -> [Text] +header logfloats mmodule = [ "{-# LANGUAGE DataKinds, NegativeLiterals #-}" - , "module Main where" + , TxT.unwords [ "module" + , case mmodule of + Just m -> pack m + Nothing -> "Main" + , "where" ] , "" - , "import Prelude hiding (product)" - , "import Language.Hakaru.Runtime.Prelude" + , if logfloats + then TxT.unlines [ "import Data.Number.LogFloat (LogFloat)" + , "import Prelude hiding (product, exp, log, (**))" + ] + else "import Prelude hiding (product)" + , if logfloats + then TxT.unlines [ "import Language.Hakaru.Runtime.LogFloatPrelude" + , "import Language.Hakaru.Runtime.LogFloatCmdLine" ] + else TxT.unlines [ "import Language.Hakaru.Runtime.Prelude" + , "import Language.Hakaru.Runtime.CmdLine" ] , "import Language.Hakaru.Types.Sing" , "import qualified System.Random.MWC as MWC" , "import Control.Monad" + , "import System.Environment (getArgs)" , "" ] -footer :: Sing (a :: Hakaru) -> [Text] -footer typ = - [ "" - , "main :: IO ()" - , "main = do" - , " g <- MWC.createSystemRandom" - , case typ of - SMeasure _ -> " forever $ run g prog" - _ -> " print prog" - ] +footer :: forall (a :: Hakaru) . Bool -> Sing a -> [Text] +footer logfloats typ = + ["","main :: IO ()" + , TxT.concat ["main = makeMain (prog :: ",toHsType typ,") =<< getArgs"]] + where toHsType :: forall (a :: Hakaru) . Sing a -> Text + toHsType SInt = "Int" + toHsType SNat = "Int" + toHsType SReal = "Double" + toHsType SProb = if logfloats then "LogFloat" else "Double" + toHsType (SArray t) = let t' = toHsType t in + TxT.concat ["(",TxT.unwords ["MayBoxVec",t',t'],")"] + toHsType (SMeasure t) = TxT.concat ["(",TxT.unwords ["Measure",toHsType t],")"] + toHsType (SFun t1 t2) = TxT.unwords [toHsType t1,"->",toHsType t2] + toHsType (SData _ + ((SKonst t1 `SEt` SKonst t2 `SEt` SDone) `SPlus` SVoid)) = + TxT.concat ["(",toHsType t1,",",toHsType t2,")"] + toHsType _ = "type" + footerWalk :: [Text] footerWalk = diff --git a/commands/Disintegrate.hs b/commands/Disintegrate.hs --- a/commands/Disintegrate.hs +++ b/commands/Disintegrate.hs @@ -1,40 +1,111 @@ -{-# LANGUAGE OverloadedStrings, PatternGuards, DataKinds, GADTs #-} +{-# LANGUAGE OverloadedStrings + , PatternGuards + , TypeOperators + , DataKinds + , GADTs + , KindSignatures + , FlexibleContexts + , RankNTypes + , CPP #-} module Main where import Language.Hakaru.Pretty.Concrete import Language.Hakaru.Syntax.TypeCheck import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST as AST import Language.Hakaru.Types.Sing +import Language.Hakaru.Types.DataKind import Language.Hakaru.Disintegrate import Language.Hakaru.Evaluation.ConstantPropagation import Language.Hakaru.Command -import Data.Text +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import Data.Monoid +import Control.Monad (when) +import qualified Data.Text as T import qualified Data.Text.IO as IO import System.IO (stderr) -import System.Environment +import Options.Applicative as O +data Options = Options { total :: Bool + , index :: Int + , program :: String } + +options :: Parser Options +options = Options + <$> switch + ( long "total" <> + short 't' <> + help "Whether to show the total number of disintegrations" ) + <*> option auto + ( long "index" <> + short 'i' <> + metavar "INDEX" <> + value 0 <> + help "The index of the desired result in the list of disintegrations (default: 0)" ) + <*> strArgument + ( metavar "PROGRAM" <> + help "File containing program to disintegrate" ) + +parseOpts :: IO Options +parseOpts = execParser $ info (helper <*> options) $ + fullDesc <> + progDesc "Disintegrate a Hakaru program" <> + header + "disintegrate: symbolic conditioning of probabilistic programs" + main :: IO () main = do - args <- getArgs - progs <- mapM readFromFile args - case progs of - [prog] -> runDisintegrate prog - _ -> IO.hPutStrLn stderr "Usage: disintegrate " + args <- parseOpts + case args of + Options t i file -> do + prog <- readFromFile file + runDisintegrate prog t i -runDisintegrate :: Text -> IO () -runDisintegrate prog = +runDisintegrate :: T.Text -> Bool -> Int -> IO () +runDisintegrate prog showTotal i = case parseAndInfer prog of - Left err -> IO.hPutStrLn stderr err - Right (TypedAST typ ast) -> - case typ of - SMeasure (SData (STyCon sym `STyApp` _ `STyApp` _) _) - | Just Refl <- jmEq1 sym sSymbol_Pair - -> case determine (disintegrate ast) of - Just ast' -> print . pretty $ constantPropagation ast' - Nothing -> IO.hPutStrLn stderr "No disintegration found" - _ -> IO.hPutStrLn stderr "Can only disintegrate a measure over pairs" + Left err -> IO.hPutStrLn stderr err + Right typedAST -> go Nil1 typedAST + where + go :: List1 Variable (xs :: [Hakaru]) + -> TypedAST (TrivialABT AST.Term) + -> IO () + go xs (TypedAST typ ast) + = case typ of + SMeasure (SData (STyCon sym `STyApp` _ `STyApp` _) _) + | Just Refl <- jmEq1 sym sSymbol_Pair + -> case disintegrate ast of + [] -> IO.hPutStrLn stderr + "No disintegration found" + rs -> when showTotal + (IO.hPutStrLn stderr.T.pack $ + "Number of disintegrations: " ++ show (length rs)) >> + lams xs (print.pretty.constantPropagation) (rs Prelude.!! i) + SFun _ b -> + caseVarSyn ast putErrorMsg $ \t -> + case t of + Lam_ :$ body :* End -> + caseBind body $ \x e -> + go (append1 xs (Cons1 x Nil1)) (TypedAST b e) + _ -> putErrorMsg ast + _ -> putErrorMsg ast + +putErrorMsg :: (Show a) => a -> IO () +putErrorMsg a = IO.hPutStrLn stderr . T.pack $ + "Can only disintegrate (functions over) measures over pairs" + -- ++ "\nGiven\n" ++ show a +-- | Use a list of variables to wrap lambdas around a given term +lams :: (ABT AST.Term abt) + => List1 Variable (xs :: [Hakaru]) + -> (forall a. abt '[] a -> IO ()) -> abt '[] a -> IO () +lams Nil1 k = k +lams (Cons1 x xs) k = lams xs (k . lam_ x) diff --git a/commands/HKC.hs b/commands/HKC.hs --- a/commands/HKC.hs +++ b/commands/HKC.hs @@ -1,11 +1,18 @@ -{-# LANGUAGE GADTs, - OverloadedStrings #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} module Main where import Language.Hakaru.Evaluation.ConstantPropagation import Language.Hakaru.Syntax.TypeCheck import Language.Hakaru.Syntax.AST.Transforms (expandTransformations) +import Language.Hakaru.Syntax.ANF (normalize) +import Language.Hakaru.Syntax.CSE (cse) +import Language.Hakaru.Syntax.Prune (prune) +import Language.Hakaru.Syntax.Uniquify (uniquify) +import Language.Hakaru.Syntax.Hoist (hoist) +import Language.Hakaru.Summary import Language.Hakaru.Command import Language.Hakaru.CodeGen.Wrapper import Language.Hakaru.CodeGen.CodeGenMonad @@ -13,6 +20,9 @@ import Language.Hakaru.CodeGen.Pretty import Control.Monad.Reader + +import Data.Monoid +import Data.Maybe (isJust) import Data.Text hiding (any,map,filter,foldr) import qualified Data.Text.IO as IO import Text.PrettyPrint (render) @@ -24,14 +34,17 @@ data Options = Options { debug :: Bool - , optimize :: Bool + , optimize :: Maybe Int + , summaryOpt :: Bool , make :: Maybe String , asFunc :: Maybe String , fileIn :: String , fileOut :: Maybe String , par :: Bool - , showWeightsOpt :: Bool + , noWeightsOpt :: Bool , showProbInLogOpt :: Bool + , garbageCollector :: Bool + -- , logProbs :: Bool } deriving Show @@ -46,9 +59,12 @@ <$> switch ( long "debug" <> short 'D' <> help "Prints Hakaru src, Hakaru AST, C AST, C src" ) - <*> switch ( long "optimize" - <> short 'O' - <> help "Performs constant folding on Hakaru AST" ) + <*> (optional $ option auto ( short 'O' + <> metavar "{0,1,2}" + <> help "perform Hakaru AST optimizations. optimization levels 0,1,2." )) + <*> switch ( long "summary" + <> short 'S' + <> help "Performs summarization optimization" ) <*> (optional $ strOption ( long "make" <> short 'm' <> help "Compiles generated C code with the compiler ARG")) @@ -62,11 +78,16 @@ <> help "output FILE")) <*> switch ( short 'j' <> help "Generates parallel programs using OpenMP directives") - <*> switch ( long "show-weights" + <*> switch ( long "no-weights" <> short 'w' - <> help "Shows the weights along with the samples in samplers") + <> help "Don't print the weights") <*> switch ( long "show-prob-log" <> help "Shows prob types as 'exp()' instead of ''") + <*> switch ( long "garbage-collector" + <> short 'g' + <> help "Use Boehm Garbage Collector") + -- <*> switch ( long "-no-log-space-probs" + -- <> help "Do not log `prob` types; WARNING this is more likely to underflow.") parseOpts :: IO Options parseOpts = execParser $ info (helper <*> options) @@ -74,27 +95,32 @@ compileHakaru :: Text -> ReaderT Options IO () compileHakaru prog = ask >>= \config -> lift $ do - case parseAndInfer prog of + prog' <- parseAndInferWithDebug (debug config) prog + case prog' of Left err -> IO.hPutStrLn stderr err Right (TypedAST typ ast) -> do - let ast' = TypedAST typ $ foldr id ast abtPasses + astS <- case summaryOpt config of + True -> summary (expandTransformations ast) + False -> return (expandTransformations ast) + let ast' = TypedAST typ $ foldr id astS (abtPasses $ optimize config) outPath = case fileOut config of (Just f) -> f Nothing -> "-" codeGen = wrapProgram ast' (asFunc config) - (PrintConfig { showWeights = showWeightsOpt config + (PrintConfig { noWeights = noWeightsOpt config , showProbInLog = showProbInLogOpt config }) - cast = CAST $ runCodeGenWith codeGen (emptyCG {parallel = par config}) + codeGenConfig = emptyCG {sharedMem = par config, managedMem = garbageCollector config} + cast = CAST $ runCodeGenWith codeGen codeGenConfig output = pack . render . pretty $ cast when (debug config) $ do putErrorLn $ hrule "Hakaru Type" putErrorLn . pack . show $ typ - when (optimize config) $ do - putErrorLn $ hrule "Hakaru AST" - putErrorLn $ pack $ show ast - putErrorLn $ hrule "Hakaru AST'" + putErrorLn $ hrule "Hakaru AST" putErrorLn $ pack $ show ast + when (isJust . optimize $ config) $ do + putErrorLn $ hrule "Hakaru AST'" + putErrorLn $ pack $ show ast' putErrorLn $ hrule "C AST" putErrorLn $ pack $ show cast putErrorLn $ hrule "Fin" @@ -102,10 +128,20 @@ Nothing -> writeToFile outPath output Just cc -> makeFile cc (fileOut config) (unpack output) config - where hrule s = concat ["\n<=======================| " - ,s," |=======================>\n"] - abtPasses = [ expandTransformations - , constantPropagation ] + where hrule s = concat [ "\n<=======================| " + , s + ," |=======================>\n"] + abtPasses Nothing = [] + abtPasses (Just 0) = [ constantPropagation ] + abtPasses (Just 1) = [ constantPropagation + , uniquify + , prune + , cse + , hoist + , uniquify ] + abtPasses (Just 2) = abtPasses (Just 1) ++ [normalize] + + abtPasses _ = error "only optimization levels are 0, 1, and 2" putErrorLn :: Text -> IO () putErrorLn = IO.hPutStrLn stderr diff --git a/commands/Hakaru.hs b/commands/Hakaru.hs --- a/commands/Hakaru.hs +++ b/commands/Hakaru.hs @@ -1,9 +1,10 @@ -{-# LANGUAGE OverloadedStrings, - PatternGuards, - DataKinds, - GADTs, - TypeOperators - #-} +{-# LANGUAGE CPP + , OverloadedStrings + , PatternGuards + , DataKinds + , GADTs + , TypeOperators + #-} module Main where @@ -17,55 +18,113 @@ import Language.Hakaru.Sample import Language.Hakaru.Pretty.Concrete -import Language.Hakaru.Command (parseAndInfer, readFromFile, Term) +import Language.Hakaru.Command ( parseAndInfer, parseAndInfer' + , readFromFile, Term + ) +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif import Control.Monad +import Data.Monoid import Data.Text import qualified Data.Text.IO as IO +import qualified Data.Vector as V +import Data.Word import System.IO (stderr) import Text.PrettyPrint (renderStyle, style, mode, Mode(LeftMode)) -import qualified System.Random.MWC as MWC -import System.Environment +import qualified Options.Applicative as O +import qualified System.Random.MWC as MWC +data Options = Options + { noWeights :: Bool + , seed :: Maybe Word32 + , transition :: Maybe String + , prog :: String } + +options :: O.Parser Options +options = Options + <$> O.switch + ( O.short 'w' <> + O.long "no-weights" <> + O.help "Don't print the weights" ) + <*> O.optional (O.option O.auto + ( O.long "seed" <> + O.help "Set random seed" <> + O.metavar "seed")) + <*> O.optional (O.strOption + ( O.long "transition-kernel" <> + O.metavar "k" <> + O.help "Use this program as transition kernel for running a markov chain")) + <*> O.strArgument + ( O.metavar "PROGRAM" <> + O.help "Hakaru program to run" ) + +parseOpts :: IO Options +parseOpts = O.execParser $ O.info (O.helper <*> options) + (O.fullDesc <> O.progDesc "Run a hakaru program") + main :: IO () main = do - args <- getArgs - g <- MWC.createSystemRandom - progs <- mapM readFromFile args - case progs of - [prog1, prog2] -> randomWalk g prog1 prog2 - [prog] -> runHakaru g prog - _ -> IO.hPutStrLn stderr - "Usage: hakaru \n\ - \ | hakaru " + args <- parseOpts + g <- case seed args of + Nothing -> MWC.createSystemRandom + Just s -> MWC.initialize (V.singleton s) + case transition args of + Nothing -> runHakaru' g (noWeights args) =<< readFromFile (prog args) + Just prog2 -> do prog' <- readFromFile (prog args) + trans <- readFromFile prog2 + randomWalk' g trans prog' -illustrate :: Sing a -> MWC.GenIO -> Value a -> IO () -illustrate (SMeasure s) g (VMeasure m) = do +-- TODO: A better needs to be found for passing weights around +illustrate :: Sing a -> Bool -> MWC.GenIO -> Value a -> IO () +illustrate (SMeasure s) weights g (VMeasure m) = do x <- m (VProb 1) g case x of - Just (samp, _) -> illustrate s g samp - Nothing -> illustrate (SMeasure s) g (VMeasure m) + Just (samp, w) -> (if weights then id else withWeight w) (illustrate s weights g samp) + Nothing -> illustrate (SMeasure s) weights g (VMeasure m) -illustrate _ _ x = render x +illustrate _ _ _ x = renderLn x +withWeight :: Value 'HProb -> IO () -> IO () +withWeight w m = render w >> putStr "\t" >> m + render :: Value a -> IO () -render = putStrLn . renderStyle style {mode = LeftMode} . prettyValue +render = putStr . renderStyle style {mode = LeftMode} . prettyValue -runHakaru :: MWC.GenIO -> Text -> IO () -runHakaru g prog = - case parseAndInfer prog of +renderLn :: Value a -> IO () +renderLn = putStrLn . renderStyle style {mode = LeftMode} . prettyValue + +-- TODO: A better needs to be found for passing weights around +runHakaru :: MWC.GenIO -> Bool -> Text -> IO () +runHakaru g weights prog' = + case parseAndInfer prog' of Left err -> IO.hPutStrLn stderr err Right (TypedAST typ ast) -> do case typ of - SMeasure _ -> forever (illustrate typ g $ run ast) - _ -> illustrate typ g $ run ast + SMeasure _ -> forever (illustrate typ weights g $ run ast) + _ -> illustrate typ weights g $ run ast where run :: Term a -> Value a run = runEvaluate . expandTransformations -randomWalk ::MWC.GenIO -> Text -> Text -> IO () +-- TODO: A better needs to be found for passing weights around +runHakaru' :: MWC.GenIO -> Bool -> Text -> IO () +runHakaru' g weights prog = do + prog' <- parseAndInfer' prog + case prog' of + Left err -> IO.hPutStrLn stderr err + Right (TypedAST typ ast) -> do + case typ of + SMeasure _ -> forever (illustrate typ weights g $ run ast) + _ -> illustrate typ weights g $ run ast + where + run :: Term a -> Value a + run = runEvaluate . expandTransformations + +randomWalk :: MWC.GenIO -> Text -> Text -> IO () randomWalk g p1 p2 = case (parseAndInfer p1, parseAndInfer p2) of (Right (TypedAST typ1 ast1), Right (TypedAST typ2 ast2)) -> @@ -84,7 +143,31 @@ chain :: Value (a ':-> b) -> Value ('HMeasure a) -> IO (Value b) chain (VLam f) (VMeasure m) = do Just (samp,_) <- m (VProb 1) g - render samp + renderLn samp + return (f samp) + +randomWalk' :: MWC.GenIO -> Text -> Text -> IO () +randomWalk' g p1 p2 = do + p1' <- parseAndInfer' p1 + p2' <- parseAndInfer' p2 + case (p1', p2') of + (Right (TypedAST typ1 ast1), Right (TypedAST typ2 ast2)) -> + -- TODO: Use better error messages for type mismatch + case (typ1, typ2) of + (SFun a (SMeasure b), SMeasure c) + | (Just Refl, Just Refl) <- (jmEq1 a b, jmEq1 b c) + -> iterateM_ (chain $ run ast1) (run ast2) + _ -> IO.hPutStrLn stderr "hakaru: programs have wrong type" + (Left err, _) -> IO.hPutStrLn stderr err + (_, Left err) -> IO.hPutStrLn stderr err + where + run :: Term a -> Value a + run = runEvaluate . expandTransformations + + chain :: Value (a ':-> b) -> Value ('HMeasure a) -> IO (Value b) + chain (VLam f) (VMeasure m) = do + Just (samp,_) <- m (VProb 1) g + renderLn samp return (f samp) -- From monad-loops diff --git a/commands/HkMaple.hs b/commands/HkMaple.hs new file mode 100644 --- /dev/null +++ b/commands/HkMaple.hs @@ -0,0 +1,128 @@ +{-# LANGUAGE CPP + , OverloadedStrings + , DataKinds + , GADTs + , RecordWildCards + , ScopedTypeVariables + #-} + +module Main where + +import Language.Hakaru.Pretty.Concrete +import Language.Hakaru.Syntax.AST.Transforms +import Language.Hakaru.Syntax.TypeCheck +import Language.Hakaru.Command (parseAndInfer, readFromFile, Term) + +import Language.Hakaru.Syntax.Rename +import Language.Hakaru.Simplify +import Language.Hakaru.Maple +import Language.Hakaru.Parser.Maple + + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import Data.Monoid ((<>), mconcat) +import Data.Text (Text, unpack, pack) +import qualified Data.Text as Text +import qualified Data.Text.IO as IO +import System.IO (stderr) +import Data.List (intercalate) +import Text.Read (readMaybe) +import Control.Exception(throw) +import qualified Options.Applicative as O +import qualified Data.Map as M + + +data Options a + = Options + { moptions :: MapleOptions String + , no_unicode :: Bool + , program :: a } + | ListCommands + + +parseKeyVal :: O.ReadM (String, String) +parseKeyVal = + O.maybeReader $ (\str -> + case map Text.strip $ Text.splitOn "," str of + [k,v] -> return (unpack k, unpack v) + _ -> Nothing) . pack + +options :: O.Parser (Options FilePath) +options = (Options + <$> (MapleOptions <$> + O.option O.str + ( O.long "command" <> + O.help ("Command to send to Maple. You may enter a prefix of the command string if " + ++"it uniquely identifies a command. " + ++"Default: Simplify") <> + O.short 'c' <> + O.value "Simplify" ) + <*> O.switch + ( O.long "debug" <> + O.help "Prints output that is sent to Maple." ) + <*> O.option O.auto + ( O.long "timelimit" <> + O.help "Set Maple to timeout in N seconds." <> + O.showDefault <> + O.value 90 <> + O.metavar "N") + <*> (M.fromList <$> + O.many (O.option parseKeyVal + ( O.long "maple-opt" <> + O.short 'm' <> + O.help ( "Extra options to send to Maple\neach options is of the form KEY=VAL\n" + ++"where KEY is a Maple name, and VAL is a Maple expression.") + )))) + <*> O.switch + ( O.long "no-unicode" <> + O.short 'u' <> + O.help "Removes unicode characters from names in the Maple output.") + <*> O.strArgument + ( O.metavar "PROGRAM" <> + O.help "Filename containing program to be simplified, or \"-\" to read from input." )) O.<|> + ( O.flag' ListCommands + ( O.long "list-commands" <> + O.help "Get list of available commands from Maple." <> + O.short 'l') ) + +parseOpts :: IO (Options FilePath) +parseOpts = O.execParser $ O.info (O.helper <*> options) + (O.fullDesc <> O.progDesc progDesc) + +progDesc :: String +progDesc = unwords + ["hk-maple: invokes a Maple command on a Hakaru program. " + ,"Given a Hakaru program in concrete syntax and a Maple-Hakaru command," + ,"typecheck the program" + ,"invoke the Maple command on the program and its type" + ,"pretty print, parse and typecheck the program resulting from Maple" + ] + +et (TypedAST t (x :: Term a)) = TypedAST t (expandTransformations x) + +main :: IO () +main = parseOpts >>= runMaple + +runMaple :: Options FilePath -> IO () +runMaple ListCommands = + listCommands >>= \cs -> putStrLn $ "Available Hakaru Maple commands:\n\t"++ intercalate ", " cs + +runMaple Options{..} = readFromFile program >>= \prog -> + case parseAndInfer prog of + Left err -> IO.hPutStrLn stderr err + Right ast -> do + TypedAST _ ast' <- sendToMaple' moptions (et ast) + print $ pretty + $ (if no_unicode then renameAST removeUnicodeChars else id) + $ ast' + +listCommands :: IO [String] +listCommands = do + let toMaple_ = "use Hakaru, NewSLO in lprint(map(curry(sprintf,`%s`),NewSLO:-Commands)) end use;" + fromMaple <- maple toMaple_ + maybe (throw $ MapleInterpreterException fromMaple toMaple_) + return + (readMaybe fromMaple) diff --git a/commands/Momiji.hs b/commands/Momiji.hs --- a/commands/Momiji.hs +++ b/commands/Momiji.hs @@ -8,25 +8,19 @@ import Language.Hakaru.Command import Data.Text -import qualified Data.Text.IO as IO -import System.IO (stderr) +import qualified Data.Text.Utf8 as U -import System.Environment +import System.IO (stderr) main :: IO () -main = do - args <- getArgs - case args of - [prog] -> IO.readFile prog >>= runPretty - [] -> IO.getContents >>= runPretty - _ -> IO.hPutStrLn stderr "Usage: momiji " +main = simpleCommand runPretty "momiji" runPretty :: Text -> IO () runPretty prog = case parseAndInfer prog of - Left err -> IO.hPutStrLn stderr err + Left err -> U.hPut stderr err Right (TypedAST typ ast) -> do - putStrLn $ Maple.pretty (expandTransformations ast) - putStrLn $ "," - putStrLn $ Maple.mapleType typ "" + U.putStrLnS $ Maple.pretty (expandTransformations ast) + U.putStrLn $ "," + U.putStrLnS $ Maple.mapleType typ "" diff --git a/commands/Normalize.hs b/commands/Normalize.hs --- a/commands/Normalize.hs +++ b/commands/Normalize.hs @@ -15,15 +15,8 @@ import qualified Data.Text.IO as IO import System.IO (stderr) -import System.Environment - main :: IO () -main = do - args <- getArgs - case args of - [prog] -> IO.readFile prog >>= runNormalize - [] -> IO.getContents >>= runNormalize - _ -> IO.hPutStrLn stderr "Usage: normalize " +main = simpleCommand runNormalize "normalize" runNormalize :: Text -> IO () runNormalize prog = either (IO.hPutStrLn stderr) print $ do diff --git a/commands/Pretty.hs b/commands/Pretty.hs --- a/commands/Pretty.hs +++ b/commands/Pretty.hs @@ -8,22 +8,15 @@ import Language.Hakaru.Command import Data.Text -import qualified Data.Text.IO as IO +import qualified Data.Text.Utf8 as IO import System.IO (stderr) -import System.Environment - main :: IO () -main = do - args <- getArgs - case args of - [prog] -> IO.readFile prog >>= runPretty - [] -> IO.getContents >>= runPretty - _ -> IO.hPutStrLn stderr "Usage: pretty " +main = simpleCommand runPretty "pretty" runPretty :: Text -> IO () runPretty prog = case parseAndInfer prog of Left err -> IO.hPutStrLn stderr err - Right (TypedAST _ ast) -> print . pretty . expandTransformations $ ast + Right (TypedAST _ ast) -> IO.print . pretty . expandTransformations $ ast diff --git a/commands/Simplify.hs b/commands/Simplify.hs deleted file mode 100644 --- a/commands/Simplify.hs +++ /dev/null @@ -1,55 +0,0 @@ -{-# LANGUAGE CPP, OverloadedStrings, DataKinds, GADTs #-} - -module Main where - -import Language.Hakaru.Pretty.Concrete -import Language.Hakaru.Syntax.AST.Transforms -import Language.Hakaru.Syntax.TypeCheck -import Language.Hakaru.Command (parseAndInfer, readFromFile, Term) - -import Language.Hakaru.Simplify - -#if __GLASGOW_HASKELL__ < 710 -import Control.Applicative (Applicative(..), (<$>)) -#endif - -import Data.Text -import qualified Data.Text.IO as IO -import System.IO (stderr) - -import qualified Options.Applicative as O - -data Options = Options - { debug :: Bool - , program :: String } - -options :: O.Parser Options -options = Options - <$> O.switch - ( O.long "debug" O.<> - O.help "Prints output that is sent to Maple" ) - <*> O.strArgument - ( O.metavar "PROGRAM" O.<> - O.help "Program to be simplified" ) - -parseOpts :: IO Options -parseOpts = O.execParser $ O.info (O.helper <*> options) - (O.fullDesc O.<> O.progDesc "Simplify a hakaru program") - -et :: Term a -> Term a -et = expandTransformations - -main :: IO () -main = do - args <- parseOpts - case args of - Options debug_ file -> do - prog <- readFromFile file - runSimplify prog debug_ - -runSimplify :: Text -> Bool -> IO () -runSimplify prog debug_ = - case parseAndInfer prog of - Left err -> IO.hPutStrLn stderr err - Right (TypedAST _ ast) -> simplifyDebug debug_ (et ast) >>= print . pretty - diff --git a/commands/Summary.hs b/commands/Summary.hs new file mode 100644 --- /dev/null +++ b/commands/Summary.hs @@ -0,0 +1,167 @@ +{-# LANGUAGE CPP, + OverloadedStrings, + PatternGuards, + DataKinds, + KindSignatures, + GADTs, + FlexibleContexts, + TypeOperators + #-} + +module Main where + +import qualified Language.Hakaru.Syntax.AST as T +import Language.Hakaru.Syntax.AST.Transforms +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.TypeCheck + +import Language.Hakaru.Types.Sing +import Language.Hakaru.Types.DataKind + +import Language.Hakaru.Pretty.Haskell +import Language.Hakaru.Command +import Language.Hakaru.Summary + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import Data.Text as TxT +import qualified Data.Text.IO as IO +import Data.Monoid ((<>)) +import System.IO (stderr) +import Text.PrettyPrint hiding ((<>)) +import qualified Options.Applicative as O +import System.FilePath + + +data Options = + Options { fileIn :: String + , fileOut :: Maybe String + , asModule :: Maybe String + , logFloatPrelude :: Bool + , optimize :: Bool + } deriving Show + +main :: IO () +main = do + opts <- parseOpts + compileHakaru opts + +parseOpts :: IO Options +parseOpts = O.execParser $ O.info (O.helper <*> options) + $ O.fullDesc <> O.progDesc "Compile Hakaru to Haskell" + +{- + +There is some redundancy in Compile.hs, Hakaru.hs, and HKC.hs in how we decide +what to run given which arguments. There may be a way to unify these. + +-} + +options :: O.Parser Options +options = Options + <$> O.strArgument (O.metavar "INPUT" <> + O.help "Program to be compiled" ) + <*> (O.optional $ + O.strOption (O.short 'o' <> + O.help "Optional output file name")) + <*> (O.optional $ + O.strOption (O.long "as-module" <> + O.short 'M' <> + O.help "creates a haskell module with this name")) + <*> O.switch (O.long "logfloat-prelude" <> + O.help "use logfloat prelude for numeric stability") + <*> O.switch (O.short 'O' <> + O.help "perform Hakaru AST optimizations") + + + +prettyProg :: (ABT T.Term abt) + => String + -> abt '[] a + -> String +prettyProg name ast = + renderStyle style + (cat [ text (name ++ " = ") + , nest 2 (pretty ast) + ]) + +compileHakaru + :: Options + -> IO () +compileHakaru opts = do + let file = fileIn opts + prog <- readFromFile file + case parseAndInfer prog of + Left err -> IO.hPutStrLn stderr err + Right (TypedAST typ ast) -> do + ast' <- (if optimize opts then optimizations else id) <$> summary (et ast) + writeHkHsToFile file (fileOut opts) . TxT.unlines $ + header (logFloatPrelude opts) (asModule opts) ++ + [ pack $ prettyProg "prog" ast' ] ++ + (case asModule opts of + Nothing -> footer typ + Just _ -> []) + where et = expandTransformations + +writeHkHsToFile :: String -> Maybe String -> Text -> IO () +writeHkHsToFile inFile moutFile content = + let outFile = case moutFile of + Nothing -> replaceFileName inFile (takeBaseName inFile) ++ ".hs" + Just x -> x + in writeToFile outFile content + +header :: Bool -> Maybe String -> [Text] +header logfloats mmodule = + [ "{-# LANGUAGE DataKinds, NegativeLiterals #-}" + , TxT.unwords [ "module" + , case mmodule of + Just m -> pack m + Nothing -> "Main" + , "where" ] + , "" + , if logfloats + then TxT.unlines [ "import Data.Number.LogFloat (LogFloat)" + , "import Prelude hiding (product, exp, log, (**))" + ] + else "import Prelude hiding (product)" + , if logfloats + then TxT.unlines [ "import Language.Hakaru.Runtime.LogFloatPrelude" + , "import Language.Hakaru.Runtime.LogFloatCmdLine" ] + else TxT.unlines [ "import Language.Hakaru.Runtime.Prelude" + , "import Language.Hakaru.Runtime.CmdLine" ] + , "import Language.Hakaru.Types.Sing" + , "import qualified System.Random.MWC as MWC" + , "import Control.Monad" + , "import System.Environment (getArgs)" + , "" + ] + +footer :: Sing (a :: Hakaru) -> [Text] +footer typ = + ["","main :: IO ()" + , TxT.concat ["main = makeMain (prog :: ",toHsType typ,") =<< getArgs"]] + where toHsType :: Sing (a :: Hakaru) -> Text + toHsType SInt = "Int" + toHsType SNat = "Int" + toHsType SReal = "Double" + toHsType SProb = "LogFloat" + toHsType (SArray t) = let t' = toHsType t in + TxT.concat ["(",TxT.unwords ["MayBoxVec",t',t'],")"] + toHsType (SMeasure t) = TxT.concat ["(",TxT.unwords ["Measure",toHsType t],")"] + toHsType (SFun t1 t2) = TxT.unwords [toHsType t1,"->",toHsType t2] + toHsType (SData _ + ((SKonst t1 `SEt` SKonst t2 `SEt` SDone) `SPlus` SVoid)) = + TxT.concat ["(",toHsType t1,",",toHsType t2,")"] + toHsType _ = "type" + +footerWalk :: [Text] +footerWalk = + [ "" + , "main :: IO ()" + , "main = do" + , " g <- MWC.createSystemRandom" + , " x <- prog2 g" + , " iterateM_ (withPrint $ flip prog1 g) x" + ] diff --git a/hakaru.cabal b/hakaru.cabal --- a/hakaru.cabal +++ b/hakaru.cabal @@ -3,7 +3,7 @@ cabal-version: >=1.16 build-type: Simple name: hakaru -version: 0.3.0 +version: 0.4.0 synopsis: A probabilistic programming language description: Hakaru is a simply-typed probabilistic programming language, designed for easy specification of probabilistic models, and inference algorithms. @@ -39,6 +39,7 @@ Language.Hakaru.Syntax.ABT, Language.Hakaru.Syntax.Variable, Language.Hakaru.Syntax.Value, + Language.Hakaru.Syntax.Reducer, Language.Hakaru.Syntax.Datum, Language.Hakaru.Syntax.DatumABT, Language.Hakaru.Syntax.DatumCase, @@ -46,23 +47,38 @@ Language.Hakaru.Types.Sing, Language.Hakaru.Types.HClasses, Language.Hakaru.Types.Coercion, + Language.Hakaru.Syntax.ANF, + Language.Hakaru.Syntax.Uniquify, + Language.Hakaru.Syntax.Unroll, Language.Hakaru.Syntax.AST, Language.Hakaru.Syntax.AST.Transforms, Language.Hakaru.Syntax.AST.Sing, Language.Hakaru.Syntax.AST.Eq, + Language.Hakaru.Syntax.Command, + Language.Hakaru.Syntax.CSE, + Language.Hakaru.Syntax.Gensym, + Language.Hakaru.Syntax.Rename, + Language.Hakaru.Syntax.Hoist, + Language.Hakaru.Syntax.Prune, Language.Hakaru.Syntax.TypeCheck, Language.Hakaru.Syntax.TypeOf, Language.Hakaru.Syntax.Prelude, Language.Hakaru.Parser.AST, Language.Hakaru.Parser.Maple, + Language.Hakaru.Parser.Import, Language.Hakaru.Parser.Parser, Language.Hakaru.Parser.SymbolResolve, Language.Hakaru.Pretty.Haskell, Language.Hakaru.Pretty.Concrete, Language.Hakaru.Pretty.Maple, Language.Hakaru.Runtime.Prelude, + Language.Hakaru.Runtime.LogFloatPrelude, + Language.Hakaru.Runtime.CmdLine, + Language.Hakaru.Runtime.LogFloatCmdLine, Language.Hakaru.Observe, + Language.Hakaru.Maple, Language.Hakaru.Simplify, + Language.Hakaru.Summary, Language.Hakaru.Sample, Language.Hakaru.Evaluation.Types, Language.Hakaru.Evaluation.Lazy, @@ -82,8 +98,10 @@ Language.Hakaru.CodeGen.Types, Language.Hakaru.CodeGen.AST, Language.Hakaru.CodeGen.Pretty, + Language.Hakaru.CodeGen.Libs, Data.Number.Nat, Data.Number.Natural + Data.Text.Utf8 other-modules: System.MapleSSH @@ -91,6 +109,7 @@ Cabal >= 1.16, ghc-prim >= 0.3 && < 0.6, transformers >= 0.3 && < 0.6, + transformers-compat >= 0.3 && < 0.6, containers >= 0.5 && < 0.6, semigroups >= 0.16, pretty >= 1.1 && < 1.2, @@ -102,13 +121,15 @@ text >= 0.11 && < 1.3, parsec >= 3.1 && < 3.2, mwc-random >= 0.13 && < 0.14, - directory >= 1.2 && < 1.3, + directory >= 1.2 && < 1.4, integration >= 0.2.0 && < 0.3.0, primitive >= 0.5 && < 0.7, process >= 1.1 && < 2.0, HUnit >= 1.2 && < 2.0, mtl >= 2.1, - filepath >= 1.1.0.2 + filepath >= 1.1.0.2, + bytestring >= 0.9, + optparse-applicative >= 0.11 && < 0.15 ---------------------------------------------------------------- Test-Suite system-testsuite @@ -117,6 +138,67 @@ Hs-Source-Dirs: haskell Default-Language: Haskell2010 GHC-Options: -Wall -fwarn-tabs + + other-modules: Data.Number.Nat + Data.Number.Natural + Data.Text.Utf8 + Language.Hakaru.Command + Language.Hakaru.Disintegrate + Language.Hakaru.Evaluation.Coalesce + Language.Hakaru.Evaluation.ConstantPropagation + Language.Hakaru.Evaluation.DisintegrationMonad + Language.Hakaru.Evaluation.EvalMonad + Language.Hakaru.Evaluation.ExpectMonad + Language.Hakaru.Evaluation.Lazy + Language.Hakaru.Evaluation.PEvalMonad + Language.Hakaru.Evaluation.Types + Language.Hakaru.Expect + Language.Hakaru.Inference + Language.Hakaru.Parser.AST + Language.Hakaru.Parser.Import + Language.Hakaru.Parser.Maple + Language.Hakaru.Parser.Parser + Language.Hakaru.Parser.SymbolResolve + Language.Hakaru.Pretty.Concrete + Language.Hakaru.Pretty.Haskell + Language.Hakaru.Pretty.Maple + Language.Hakaru.Sample + Language.Hakaru.Simplify + Language.Hakaru.Syntax.ABT + Language.Hakaru.Syntax.ANF + Language.Hakaru.Syntax.AST + Language.Hakaru.Syntax.AST.Eq + Language.Hakaru.Syntax.AST.Sing + Language.Hakaru.Syntax.CSE + Language.Hakaru.Syntax.Datum + Language.Hakaru.Syntax.DatumABT + Language.Hakaru.Syntax.DatumCase + Language.Hakaru.Syntax.Gensym + Language.Hakaru.Syntax.Hoist + Language.Hakaru.Syntax.IClasses + Language.Hakaru.Syntax.Prelude + Language.Hakaru.Syntax.Prune + Language.Hakaru.Syntax.Reducer + Language.Hakaru.Syntax.TypeCheck + Language.Hakaru.Syntax.TypeOf + Language.Hakaru.Syntax.Uniquify + Language.Hakaru.Syntax.Unroll + Language.Hakaru.Syntax.Value + Language.Hakaru.Syntax.Variable + Language.Hakaru.Types.Coercion + Language.Hakaru.Types.DataKind + Language.Hakaru.Types.HClasses + Language.Hakaru.Types.Sing + System.MapleSSH + Tests.ASTTransforms + Tests.Disintegrate + Tests.Models + Tests.Parser + Tests.RoundTrip + Tests.Sample + Tests.Simplify + Tests.TestTools + Tests.TypeCheck Build-Depends: base >= 4.6 && < 5.0, Cabal >= 1.16, @@ -138,7 +220,10 @@ process >= 1.1 && < 2.0, mtl >= 2.1, vector >= 0.10, - text >= 0.11 && < 1.3 + text >= 0.11 && < 1.3, + bytestring >= 0.9, + directory >= 1.2 && < 1.4, + optparse-applicative >= 0.11 && < 0.15 ---------------------------------------------------------------- Executable hakaru @@ -147,11 +232,13 @@ Default-Language: Haskell2010 GHC-Options: -O2 -Wall -fwarn-tabs - build-depends: base >= 4.7 && < 5.0, - mwc-random >= 0.13 && < 0.14, - text >= 0.11 && < 1.3, - pretty >= 1.1 && < 1.2, - hakaru >= 0.3 + build-depends: base >= 4.7 && < 5.0, + mwc-random >= 0.13 && < 0.14, + text >= 0.11 && < 1.3, + pretty >= 1.1 && < 1.2, + vector >= 0.10, + optparse-applicative >= 0.11 && < 0.15, + hakaru >= 0.3 ---------------------------------------------------------------- Executable compile @@ -160,16 +247,17 @@ Default-Language: Haskell2010 GHC-Options: -O2 -Wall -fwarn-tabs - build-depends: base >= 4.7 && < 5.0, - mwc-random >= 0.13 && < 0.14, - text >= 0.11 && < 1.3, - pretty >= 1.1 && < 1.2, - filepath >= 1.3, - hakaru >= 0.3 + build-depends: base >= 4.7 && < 5.0, + mwc-random >= 0.13 && < 0.14, + text >= 0.11 && < 1.3, + pretty >= 1.1 && < 1.2, + filepath >= 1.3, + optparse-applicative >= 0.11 && < 0.15, + hakaru >= 0.3 ---------------------------------------------------------------- -Executable simplify - Main-is: Simplify.hs +Executable summary + Main-is: Summary.hs Hs-Source-Dirs: commands Default-Language: Haskell2010 GHC-Options: -O2 -Wall -fwarn-tabs @@ -178,10 +266,26 @@ mwc-random >= 0.13 && < 0.14, text >= 0.11 && < 1.3, pretty >= 1.1 && < 1.2, - optparse-applicative >= 0.11 && < 0.13, + filepath >= 1.3, + optparse-applicative >= 0.11 && < 0.15, hakaru >= 0.3 ---------------------------------------------------------------- +Executable hk-maple + Main-is: HkMaple.hs + Hs-Source-Dirs: commands + Default-Language: Haskell2010 + GHC-Options: -O2 -Wall -fwarn-tabs + + build-depends: base >= 4.7 && < 5.0, + mwc-random >= 0.13 && < 0.14, + text >= 0.11 && < 1.3, + pretty >= 1.1 && < 1.2, + optparse-applicative >= 0.13 && < 0.15, + containers >= 0.5 && < 0.6, + hakaru >= 0.3 + +---------------------------------------------------------------- Executable density Main-is: Density.hs Hs-Source-Dirs: commands @@ -193,7 +297,7 @@ text >= 0.11 && < 1.3, pretty >= 1.1 && < 1.2, hakaru >= 0.3 - + ---------------------------------------------------------------- Executable disintegrate Main-is: Disintegrate.hs @@ -201,11 +305,12 @@ Default-Language: Haskell2010 GHC-Options: -O2 -Wall -fwarn-tabs - build-depends: base >= 4.7 && < 5.0, - mwc-random >= 0.13 && < 0.14, - text >= 0.11 && < 1.3, - pretty >= 1.1 && < 1.2, - hakaru >= 0.3 + build-depends: base >= 4.7 && < 5.0, + mwc-random >= 0.13 && < 0.14, + text >= 0.11 && < 1.3, + pretty >= 1.1 && < 1.2, + optparse-applicative >= 0.11 && < 0.15, + hakaru >= 0.3 ---------------------------------------------------------------- Executable pretty @@ -256,11 +361,11 @@ containers >= 0.5 && < 0.6, text >= 0.11 && < 1.3, mtl >= 2.1, - optparse-applicative >= 0.11 && < 0.13, + optparse-applicative >= 0.11 && < 0.15, pretty >= 1.1 && < 1.2, process >= 1.1 && < 2.0, - hakaru >= 0.3, - semigroups >= 0.16 + semigroups >= 0.16, + hakaru >= 0.3 ---------------------------------------------------------------- diff --git a/haskell/Data/Text/Utf8.hs b/haskell/Data/Text/Utf8.hs new file mode 100644 --- /dev/null +++ b/haskell/Data/Text/Utf8.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE CPP, OverloadedStrings #-} + +module Data.Text.Utf8 where + +import Prelude hiding (putStr, putStrLn) +import Control.Applicative (liftA) + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import qualified Data.ByteString.Char8 as BIO +import qualified Data.Text as Text +import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import System.IO (Handle) +import Data.Monoid ((<>)) + +readFile :: FilePath -> IO Text.Text +readFile f = decodeUtf8 <$> (BIO.readFile f) + +getContents :: IO Text.Text +getContents = decodeUtf8 <$> BIO.getContents + +putStr :: Text.Text -> IO () +putStr = BIO.putStr . encodeUtf8 + +putStrLn :: Text.Text -> IO () +putStrLn = BIO.putStrLn . encodeUtf8 + +writeFile :: FilePath -> Text.Text -> IO () +writeFile f x = BIO.writeFile f (encodeUtf8 x) + +hPut :: Handle -> Text.Text -> IO () +hPut h x = BIO.hPut h (encodeUtf8 x) + +hPutStrLn :: Handle -> Text.Text -> IO () +hPutStrLn h x = BIO.hPutStrLn h (encodeUtf8 x) + +putStrS :: String -> IO () +putStrS = putStr . Text.pack + +putStrLnS :: String -> IO () +putStrLnS = putStrLn . Text.pack + +print :: Show a => a -> IO () +print = BIO.putStrLn . encodeUtf8 . Text.pack . show + diff --git a/haskell/Language/Hakaru/CodeGen/AST.hs b/haskell/Language/Hakaru/CodeGen/AST.hs --- a/haskell/Language/Hakaru/CodeGen/AST.hs +++ b/haskell/Language/Hakaru/CodeGen/AST.hs @@ -19,22 +19,25 @@ -- declaration constructors , CDecl(..), CDeclr(..), CDeclSpec(..), CStorageSpec(..), CTypeQual(..) - , CDirectDeclr(..), CTypeSpec(..), CSUSpec(..), CSUTag(..), CEnum(..) - , CInit(..), CPartDesig(..), CFunSpec(..), CPtrDeclr(..) + , CDirectDeclr(..), CTypeSpec(..), CTypeName(..), CSUSpec(..), CSUTag(..) + , CEnum(..), CInit(..), CPartDesig(..), CFunSpec(..), CPtrDeclr(..) -- statements and expression constructors , CStat(..), CCompoundBlockItem(..), CExpr(..), CConst(..), CUnaryOp(..) , CBinaryOp(..), CAssignOp(..) -- infix and smart constructors - , (.>.),(.<.),(.==.),(.||.),(.&&.),(.*.),(./.),(.-.),(.+.),(.=.),(.+=.),(.*=.) - , (.>=.),(.<=.) + , (.>.),(.<.),(.==.),(.!=.),(.||.),(.&&.),(.*.),(./.),(.-.),(.+.),(.=.),(.+=.) + , (.*=.),(.>=.),(.<=.),(...),(.->.) , seqCStat - , indirect, address, intE, floatE, stringE, mkUnary - , exp, expm1, log, log1p, sqrt, rand, infinityE, negInfinityE, printfE + , indirect, address, index, intE, charE, floatE, stringE, mkCallE, mkUnaryE + , nullE + + -- util + , cNameStream ) where -import Prelude hiding (exp,log,sqrt) +import Control.Monad (mplus) -------------------------------------------------------------------------------- @@ -134,6 +137,13 @@ | CEnumType CEnum deriving (Show, Eq, Ord) +-- CTypeName is necessary for cast operations, see C99 pp81 and pp122 +-- For now, we only need to use these casts for malloc, so this is +-- incomplete with respect to C99 +data CTypeName + = CTypeName [CTypeSpec] Bool + deriving (Show, Eq, Ord) + data CSUSpec = CSUSpec CSUTag (Maybe Ident) [CDecl] deriving (Show, Eq, Ord) @@ -153,20 +163,22 @@ {- Declarators give us labels to point at and describe the level of indirection. between a label and the underlieing memory + + this is incomplete, see c99 reference p115 -} data CDeclr - = CDeclr (Maybe CPtrDeclr) [CDirectDeclr] + = CDeclr (Maybe CPtrDeclr) CDirectDeclr deriving (Show, Eq, Ord) data CPtrDeclr = CPtrDeclr [CTypeQual] deriving (Show, Eq, Ord) --- this is incomplete data CDirectDeclr = CDDeclrIdent Ident - | CDDeclrArr CDirectDeclr CExpr + | CDDeclrArr CDirectDeclr (Maybe CExpr) | CDDeclrFun CDirectDeclr [CTypeSpec] + | CDDeclrRec CDeclr deriving (Show, Eq, Ord) ------------------ @@ -232,10 +244,10 @@ | CAssign CAssignOp CExpr CExpr | CCond CExpr CExpr CExpr | CBinary CBinaryOp CExpr CExpr - | CCast CDecl CExpr + | CCast CTypeName CExpr | CUnary CUnaryOp CExpr | CSizeOfExpr CExpr - | CSizeOfType CDecl + | CSizeOfType CTypeName | CIndex CExpr CExpr | CCall CExpr [CExpr] | CMember CExpr Ident Bool @@ -315,13 +327,14 @@ seqCStat :: [CStat] -> CStat seqCStat = CCompound . fmap CBlockStat -(.<.),(.>.),(.==.),(.||.),(.&&.),(.*.),(./.),(.-.),(.+.),(.=.),(.+=.),(.*=.),(.<=.),(.>=.) +(.<.),(.>.),(.==.),(.!=.),(.||.),(.&&.),(.*.),(./.),(.-.),(.+.),(.=.),(.+=.),(.*=.),(.<=.),(.>=.) :: CExpr -> CExpr -> CExpr a .<. b = CBinary CLeOp a b a .>. b = CBinary CGrOp a b a .==. b = CBinary CEqOp a b +a .!=. b = CBinary CNeqOp a b a .||. b = CBinary CLorOp a b -a .&&. b = CBinary CAndOp a b +a .&&. b = CBinary CLndOp a b a .*. b = CBinary CMulOp a b a ./. b = CBinary CDivOp a b a .-. b = CBinary CSubOp a b @@ -332,57 +345,44 @@ a .+=. b = CAssign CAddAssOp a b a .*=. b = CAssign CMulAssOp a b + indirect, address :: CExpr -> CExpr indirect = CUnary CIndOp address = CUnary CAdrOp +index :: CExpr -> CExpr -> CExpr +index = CIndex + +(...),(.->.) :: CExpr -> String -> CExpr +i ... n = CMember i (Ident n) True +i .->. n = CMember i (Ident n) False + intE :: Integer -> CExpr intE = CConstant . CIntConst floatE :: Float -> CExpr floatE = CConstant . CFloatConst +charE :: Char -> CExpr +charE = CConstant . CCharConst + stringE :: String -> CExpr stringE = CConstant . CStringConst -mkUnary :: String -> CExpr -> CExpr -mkUnary s = CCall (CVar . Ident $ s) . (:[]) - --------------------------------------------------------------------------------- --- Lib C -- --------------------------------------------------------------------------------- -{- - Here we have calls to a very small subset of functionality provided by libc. - In the future, we should have a standard way to add in bindings to C - libraries. Easily generating code for existing C libraries is one of the key - design goals of pedantic-c --} - ------------- --- math.h -- ------------- - -exp,expm1,log,log1p,sqrt :: CExpr -> CExpr -exp = mkUnary "exp" -expm1 = mkUnary "expm1" -log = mkUnary "log" -log1p = mkUnary "log1p" -sqrt = mkUnary "sqrt" - -infinityE,negInfinityE :: CExpr -infinityE = (intE 1) ./. (intE 0) -negInfinityE = log (intE 0) +mkCallE :: String -> [CExpr] -> CExpr +mkCallE s = CCall (CVar . Ident $ s) --------------- --- stdlib.h -- --------------- +mkUnaryE :: String -> CExpr -> CExpr +mkUnaryE s a = mkCallE s [a] -rand :: CExpr -rand = CCall (CVar . Ident $ "rand") [] +nullE :: CExpr +nullE = CVar . Ident $ "NULL" --------------- --- stdlio.h -- --------------- +-------------------------------------------------------------------------------- -printfE :: [CExpr] -> CExpr -printfE = CCall (CVar . Ident $ "printf") +cNameStream :: [String] +cNameStream = filter (\n -> not $ elem (head n) ['0'..'9']) names + where base :: [Char] + base = ['0'..'9'] ++ ['a'..'z'] + names = [[x] | x <- base] `mplus` (do n <- names + [n++[x] | x <- base]) diff --git a/haskell/Language/Hakaru/CodeGen/CodeGenMonad.hs b/haskell/Language/Hakaru/CodeGen/CodeGenMonad.hs --- a/haskell/Language/Hakaru/CodeGen/CodeGenMonad.hs +++ b/haskell/Language/Hakaru/CodeGen/CodeGenMonad.hs @@ -33,7 +33,6 @@ , runCodeGenBlock , runCodeGenWith , emptyCG - , suffixes -- codegen effects , declare @@ -42,25 +41,32 @@ , putStat , putExprStat , extDeclare - , defineFunction + , extDeclareTypes + , funCG , isParallel , mkParallel , mkSequential - , reserveName + , reserveIdent , genIdent , genIdent' -- Hakaru specific , createIdent + , createIdent' , lookupIdent -- control mechanisms + , ifCG , whileCG , doWhileCG , forCG , reductionCG + , codeBlockCG + + -- memory + , putMallocStat ) where import Control.Monad.State.Strict @@ -76,6 +82,7 @@ import Language.Hakaru.CodeGen.Types import Language.Hakaru.CodeGen.AST import Language.Hakaru.CodeGen.Pretty +import Language.Hakaru.CodeGen.Libs import Data.Number.Nat (fromNat) import qualified Data.IntMap.Strict as IM @@ -84,13 +91,6 @@ import Text.PrettyPrint (render) -suffixes :: [String] -suffixes = filter (\n -> not $ elem (head n) ['0'..'9']) names - where base :: [Char] - base = ['0'..'9'] ++ ['a'..'z'] - names = [[x] | x <- base] `mplus` (do n <- names - [n++[x] | x <- base]) - -- CG after "codegen", holds the state of a codegen computation data CG = CG { freshNames :: [String] -- ^ fresh names for code generations , reservedNames :: S.Set String -- ^ reserve names during code generations @@ -98,11 +98,14 @@ , declarations :: [CDecl] -- ^ declarations in local block , statements :: [CStat] -- ^ statements can include assignments as well as other side-effects , varEnv :: Env -- ^ mapping between Hakaru vars and codegeneration vars - , parallel :: Bool -- ^ openMP supported block + , managedMem :: Bool -- ^ garbage collected block + , sharedMem :: Bool -- ^ shared memory supported block + , distributed :: Bool -- ^ distributed supported block + , logProbs :: Bool -- ^ true by default, but might not matter in some situations } emptyCG :: CG -emptyCG = CG suffixes mempty mempty [] [] emptyEnv False +emptyCG = CG cNameStream mempty mempty [] [] emptyEnv False False False True type CodeGen = State CG @@ -130,39 +133,28 @@ -------------------------------------------------------------------------------- isParallel :: CodeGen Bool -isParallel = parallel <$> get +isParallel = sharedMem <$> get mkParallel :: CodeGen () mkParallel = do cg <- get - put (cg { parallel = True } ) + put (cg { sharedMem = True } ) mkSequential :: CodeGen () mkSequential = do cg <- get - put (cg { parallel = False } ) + put (cg { sharedMem = False } ) -------------------------------------------------------------------------------- -reserveName :: String -> CodeGen () -reserveName s = +reserveIdent :: String -> CodeGen Ident +reserveIdent s = do get >>= \cg -> put $ cg { reservedNames = s `S.insert` reservedNames cg } + return (Ident s) genIdent :: CodeGen Ident -genIdent = - do cg <- get - let (freshNs,name) = pullName (freshNames cg) (reservedNames cg) - put $ cg { freshNames = freshNs } - return $ Ident name - where pullName :: [String] -> S.Set String -> ([String],String) - pullName (n:names) reserved = - let name = "_" ++ n in - if S.member name reserved - then let (names',out) = pullName names reserved - in (n:names',out) - else (names,name) - pullName _ _ = error "should not happen, names is infinite" +genIdent = genIdent' "" genIdent' :: String -> CodeGen Ident genIdent' s = @@ -182,13 +174,21 @@ createIdent :: Variable (a :: Hakaru) -> CodeGen Ident -createIdent var@(Variable name _ _) = +createIdent = createIdent' "" + +createIdent' :: String -> Variable (a :: Hakaru) -> CodeGen Ident +createIdent' s var@(Variable name _ _) = do !cg <- get - let ident = Ident $ (T.unpack name) ++ "_" ++ (head $ freshNames cg) + let ident = Ident $ concat [concatMap toAscii . T.unpack $ name + ,"_",s,"_",head $ freshNames cg ] env' = updateEnv var ident (varEnv cg) put $! cg { freshNames = tail $ freshNames cg , varEnv = env' } return ident + where toAscii c = let num = fromEnum c in + if num < 48 || num > 122 + then "u" ++ (show num) + else [c] lookupIdent :: Variable (a :: Hakaru) -> CodeGen Ident lookupIdent var = @@ -202,22 +202,48 @@ -- code in the CodeGenMonad while literal types SReal, SInt, SNat, and SProb -- do not declare :: Sing (a :: Hakaru) -> Ident -> CodeGen () -declare SInt = declare' . typeDeclaration SInt -declare SNat = declare' . typeDeclaration SNat -declare SProb = declare' . typeDeclaration SProb -declare SReal = declare' . typeDeclaration SReal -declare (SMeasure (SArray t)) = \i -> do extDeclare $ arrayStruct t - extDeclare $ mdataStruct t - declare' $ mdataDeclaration (SArray t) i -declare (SMeasure t) = \i -> do extDeclare $ mdataStruct t - declare' $ mdataDeclaration t i -declare (SArray t) = \i -> do extDeclare $ arrayStruct t - declare' $ arrayDeclaration t i -declare d@(SData _ _) = \i -> do extDeclare $ datumStruct d - declare' $ datumDeclaration d i -declare (SFun _ _) = \_ -> return () -- function definitions handeled in flatten +declare SInt = declare' . typeDeclaration SInt +declare SNat = declare' . typeDeclaration SNat +declare SProb = declare' . typeDeclaration SProb +declare SReal = declare' . typeDeclaration SReal +declare m@(SMeasure t) = \i -> + extDeclareTypes m >> (declare' $ mdataDeclaration t i) +declare a@(SArray t) = \i -> + extDeclareTypes a >> (declare' $ arrayDeclaration t i) +declare d@(SData _ _) = \i -> + extDeclareTypes d >> (declare' $ datumDeclaration d i) + +declare f@(SFun _ _) = \_ -> + extDeclareTypes f >> return () + -- this currently avoids declaration if the type is a lambda, this is hacky + +-- | for types that contain subtypes we need to recursively traverse them and +-- build up a list of external type declarations. +-- For example: Measure (Array Nat) will need to have structures for arrays +-- declared before the top level type +extDeclareTypes :: Sing (a :: Hakaru) -> CodeGen () +extDeclareTypes SInt = return () +extDeclareTypes SNat = return () +extDeclareTypes SReal = return () +extDeclareTypes SProb = return () +extDeclareTypes (SMeasure i) = extDeclareTypes i >> extDeclare (mdataStruct i) +extDeclareTypes (SArray i) = extDeclareTypes i >> extDeclare (arrayStruct i) +extDeclareTypes (SFun x y) = extDeclareTypes x >> extDeclareTypes y +extDeclareTypes d@(SData _ i) = extDeclDatum i >> extDeclare (datumStruct d) + where extDeclDatum :: Sing (a :: [[HakaruFun]]) -> CodeGen () + extDeclDatum SVoid = return () + extDeclDatum (SPlus p s) = extDeclDatum s >> datumProdTypes p + + datumProdTypes :: Sing (a :: [HakaruFun]) -> CodeGen () + datumProdTypes SDone = return () + datumProdTypes (SEt x p) = datumProdTypes p >> datumPrimTypes x + + datumPrimTypes :: Sing (a :: HakaruFun) -> CodeGen () + datumPrimTypes SIdent = return () + datumPrimTypes (SKonst s) = extDeclareTypes s + declare' :: CDecl -> CodeGen () declare' d = do cg <- get put $ cg { declarations = d:(declarations cg) } @@ -241,39 +267,21 @@ else d:extds put $ cg { extDecls = extds' } -defineFunction :: Sing (a :: Hakaru) -> Ident -> [CDecl] -> CodeGen () -> CodeGen () -defineFunction typ ident args mbody = - do cg <- get - mbody - !cg' <- get - let decls = reverse . declarations $ cg' - stmts = reverse . statements $ cg' - def :: Sing (a :: Hakaru) -> CFunDef - def SInt = functionDef SInt ident args decls stmts - def SNat = functionDef SNat ident args decls stmts - def SProb = functionDef SProb ident args decls stmts - def SReal = functionDef SReal ident args decls stmts - def (SMeasure t) = functionDef (SMeasure t) ident args decls stmts - def t = error $ "TODO: defined function of type: " ++ show t - - -- reset local statements and declarations - put $ cg' { statements = statements cg - , declarations = declarations cg } - extDeclare . CFunDefExt $ def typ - funCG :: CTypeSpec -> Ident -> [CDecl] -> CodeGen () -> CodeGen () -funCG ts ident args mbody = +funCG ts ident args m = do cg <- get - mbody - !cg' <- get + let (_,cg') = runState m $ cg { statements = [] + , declarations = [] + , freshNames = cNameStream } let decls = reverse . declarations $ cg' stmts = reverse . statements $ cg' -- reset local statements and declarations put $ cg' { statements = statements cg - , declarations = declarations cg } + , declarations = declarations cg + , freshNames = freshNames cg } extDeclare . CFunDefExt $ CFunDef [CTypeSpec ts] - (CDeclr Nothing [ CDDeclrIdent ident ]) + (CDeclr Nothing (CDDeclrIdent ident)) args (CCompound ((fmap CBlockDecl decls) ++ (fmap CBlockStat stmts))) @@ -301,19 +309,44 @@ ---------------------------------------------------------------- -- Control Flow +ifCG :: CExpr -> CodeGen () -> CodeGen () -> CodeGen () +ifCG bE m1 m2 = + do cg <- get + let (_,cg') = runState m1 $ cg { statements = [] + , declarations = [] } + (_,cg'') = runState m2 $ cg' { statements = [] + , declarations = [] } + thnBlock = (fmap CBlockDecl (reverse $ declarations cg')) + ++ (fmap CBlockStat (reverse $ statements cg')) + elsBlock = (fmap CBlockDecl (reverse $ declarations cg'') + ++ (fmap CBlockStat (reverse $ statements cg''))) + put $ cg'' { statements = statements cg + , declarations = declarations cg } + putStat $ CIf bE + (CCompound thnBlock) + (case elsBlock of + [] -> Nothing + _ -> Just . CCompound $ elsBlock) + +whileCG' :: Bool -> CExpr -> CodeGen () -> CodeGen () +whileCG' isDoWhile bE m = + do cg <- get + let (_,cg') = runState m $ cg { statements = [] + , declarations = [] } + put $ cg' { statements = statements cg + , declarations = declarations cg } + putStat $ CWhile bE + (CCompound $ (fmap CBlockDecl (reverse $ declarations cg') + ++ (fmap CBlockStat (reverse $ statements cg')))) + isDoWhile whileCG :: CExpr -> CodeGen () -> CodeGen () -whileCG bE m = - let (_,_,stmts) = runCodeGen m - in putStat $ CWhile bE (CCompound $ fmap CBlockStat stmts) False +whileCG = whileCG' False doWhileCG :: CExpr -> CodeGen () -> CodeGen () -doWhileCG bE m = - let (_,_,stmts) = runCodeGen m - in putStat $ CWhile bE (CCompound $ fmap CBlockStat stmts) True +doWhileCG = whileCG' True -- forCG and reductionCG both create C for loops, their difference lies in the -- parallel code they generate - forCG :: CExpr -> CExpr @@ -323,9 +356,11 @@ forCG iter cond inc body = do cg <- get let (_,cg') = runState body $ cg { statements = [] - , declarations = [] } - put $ cg' { statements = statements cg - , declarations = declarations cg } + , declarations = [] + , sharedMem = False } + put $ cg' { statements = statements cg + , declarations = declarations cg + , sharedMem = sharedMem cg } -- only use pragmas at the top level par <- isParallel when par . putStat . CPPStat . PPPragma $ ["omp","parallel","for"] @@ -335,8 +370,12 @@ (CCompound $ (fmap CBlockDecl (reverse $ declarations cg') ++ (fmap CBlockStat (reverse $ statements cg')))) +{- +The operation for a reduction is either a builtin binary op, or must be +specified +-} reductionCG - :: CBinaryOp + :: Either CBinaryOp (CExpr -> CExpr -> CExpr) -> Ident -> CExpr -> CExpr @@ -346,19 +385,63 @@ reductionCG op acc iter cond inc body = do cg <- get let (_,cg') = runState body $ cg { statements = [] - , declarations = [] } - put $ cg' { statements = statements cg - , declarations = declarations cg } + , declarations = [] + , sharedMem = False } -- only use pragmas at the top level + put $ cg' { statements = statements cg + , declarations = declarations cg + , sharedMem = sharedMem cg } par <- isParallel - when par . putStat . CPPStat . PPPragma - $ ["omp","parallel","for" - , concat ["reduction(" - ,render . pretty $ op - ,":" - ,render . pretty $ acc - ,")"]] + when par $ + case op of + Left binop -> putStat . CPPStat . PPPragma $ + [ "omp","parallel","for","reduction(" + , render . pretty $ binop + , ":" + , render . pretty $ acc + ,")"] + -- INCOMPLETE + Right red -> do (Ident redName) <- genIdent' "red" + let declRedPragma = [ "omp","declare","reduction(" + , redName,":",undefined,":" + , render . pretty $ + red (CVar . Ident $ "omp_in") + (CVar . Ident $ "omp_out") + , ")"] + redPragma = [ "omp","parallel","for","reduction(" + , redName + , ":" + , render . pretty $ acc + ,")"] + putStat . CPPStat . PPPragma $ declRedPragma + putStat . CPPStat . PPPragma $ redPragma putStat $ CFor (Just iter) (Just cond) (Just inc) (CCompound $ (fmap CBlockDecl (reverse $ declarations cg') ++ (fmap CBlockStat (reverse $ statements cg')))) + + +-- not control flow, but like these it creates a block with local variables +codeBlockCG :: CodeGen () -> CodeGen () +codeBlockCG body = + do cg <- get + let (_,cg') = runState body $ cg { statements = [] + , declarations = [] } + put $ cg' { statements = statements cg + , declarations = declarations cg } + putStat $ (CCompound $ (fmap CBlockDecl (reverse $ declarations cg') + ++ (fmap CBlockStat (reverse $ statements cg')))) + + + +-------------------------------------------------------------------------------- +-- ^ Takes a cexpression for the location and size and a hakaru type, and +-- generates a statement for allocating the memory +putMallocStat :: CExpr -> CExpr -> Sing (a :: Hakaru) -> CodeGen () +putMallocStat loc size typ = do + isManagedMem <- managedMem <$> get + let malloc' = if isManagedMem then gcMalloc else mallocE + typ' = buildType typ + putExprStat $ loc + .=. ( CCast (CTypeName typ' True) + $ malloc' (size .*. (CSizeOfType (CTypeName typ' False)))) diff --git a/haskell/Language/Hakaru/CodeGen/Flatten.hs b/haskell/Language/Hakaru/CodeGen/Flatten.hs --- a/haskell/Language/Hakaru/CodeGen/Flatten.hs +++ b/haskell/Language/Hakaru/CodeGen/Flatten.hs @@ -27,1218 +27,1538 @@ module Language.Hakaru.CodeGen.Flatten ( flattenABT , flattenVar - , flattenTerm ) - where - -import Language.Hakaru.CodeGen.CodeGenMonad -import Language.Hakaru.CodeGen.AST -import Language.Hakaru.CodeGen.Types - -import Language.Hakaru.Syntax.AST -import Language.Hakaru.Syntax.ABT -import Language.Hakaru.Syntax.TypeOf (typeOf) -import Language.Hakaru.Syntax.Datum hiding (Ident) -import qualified Language.Hakaru.Syntax.Prelude as HKP -import Language.Hakaru.Types.DataKind -import Language.Hakaru.Types.HClasses -import Language.Hakaru.Types.Coercion -import Language.Hakaru.Types.Sing - -import Control.Monad.State.Strict -import Control.Monad (replicateM) -import Data.Number.Natural -import Data.Ratio -import qualified Data.List.NonEmpty as NE -import qualified Data.Sequence as S -import qualified Data.Foldable as F -import qualified Data.Traversable as T - - -#if __GLASGOW_HASKELL__ < 710 -import Data.Functor -#endif - -import Prelude hiding (log,exp,sqrt) - - -opComment :: String -> CStat -opComment opStr = CComment $ concat [space," ",opStr," ",space] - where size = (50 - (length opStr)) `div` 2 - 8 - space = replicate size '-' - --------------------------------------------------------------------------------- --- Top Level -- --------------------------------------------------------------------------------- -{- - -flattening an ABT will produce a continuation that takes a CExpr representing -a location where the value of the ABT should be stored. Return type of the -the continuation is CodeGen Bool, where the computed bool is whether or not -there is a Reject inside the ABT. Therefore it is only needed when computing -mochastic values - --} - -flattenABT - :: ABT Term abt - => abt '[] a - -> (CExpr -> CodeGen ()) -flattenABT abt = caseVarSyn abt flattenVar flattenTerm - --- note that variables will find their values in the state of the CodeGen monad -flattenVar - :: Variable (a :: Hakaru) - -> (CExpr -> CodeGen ()) -flattenVar v = \loc -> - do v' <- CVar <$> lookupIdent v - putStat . CExpr . Just $ loc .=. v' - -flattenTerm - :: ABT Term abt - => Term abt a - -> (CExpr -> CodeGen ()) -flattenTerm (NaryOp_ t s) = flattenNAryOp t s -flattenTerm (Literal_ x) = flattenLit x -flattenTerm (Empty_ _) = error "TODO: flattenTerm Empty" - -flattenTerm (Datum_ d) = flattenDatum d -flattenTerm (Case_ c bs) = flattenCase c bs - -flattenTerm (Array_ s e) = flattenArray s e - --- SCon can contain mochastic terms -flattenTerm (x :$ ys) = flattenSCon x ys - ---------------------- --- Mochastic Terms -- ---------------------- -flattenTerm (Reject_ _) = \loc -> putExprStat (mdataPtrWeight loc .=. (intE 0)) -- fail to draw a sample -flattenTerm (Superpose_ wes) = flattenSuperpose wes - - - - - --------------------------------------------------------------------------------- --- SCon -- --------------------------------------------------------------------------------- - -flattenSCon - :: ( ABT Term abt ) - => SCon args a - -> SArgs abt args - -> (CExpr -> CodeGen ()) - -flattenSCon Let_ = - \(expr :* body :* End) -> - \loc -> do - caseBind body $ \v@(Variable _ _ typ) body'-> - do ident <- createIdent v - declare typ ident - flattenABT expr (CVar ident) - flattenABT body' loc - --- Lambdas produce functions and then return a function pointer -flattenSCon Lam_ = undefined - -- \(body :* End) -> - -- \loc -> - -- do coalesceLambda body $ \vars body' -> - -- let varMs = foldMap11 (\v -> [mkVarDecl v =<< createIdent v]) vars - -- in do funcId <- genIdent' "fn" - -- argDecls <- sequence varMs - - -- cg <- get - -- let m = putStat . CReturn . Just =<< flattenABT body' - -- (_,cg') = runState m $ cg { statements = [] - -- , declarations = [] } - -- put $ cg' { statements = statements cg - -- , declarations = declarations cg } - - -- extDeclare . CFunDefExt $ functionDef (typeOf body') - -- funcId - -- argDecls - -- (reverse $ declarations cg') - -- (reverse $ statements cg') - -- -- do at top level - -- where coalesceLambda - -- :: ( ABT Term abt ) - -- => abt '[x] a - -- -> (forall (ys :: [Hakaru]) b. List1 Variable ys -> abt '[] b -> r) - -- -> r - -- coalesceLambda abt k = - -- caseBind abt $ \v abt' -> - -- caseVarSyn abt' (const (k (Cons1 v Nil1) abt')) $ \term -> - -- case term of - -- (Lam_ :$ body :* End) -> - -- coalesceLambda body $ \vars abt'' -> k (Cons1 v vars) abt'' - -- _ -> k (Cons1 v Nil1) abt' - - -- mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CodeGen CDecl - -- mkVarDecl (Variable _ _ SInt) = return . typeDeclaration SInt - -- mkVarDecl (Variable _ _ SNat) = return . typeDeclaration SNat - -- mkVarDecl (Variable _ _ SProb) = return . typeDeclaration SProb - -- mkVarDecl (Variable _ _ SReal) = return . typeDeclaration SReal - -- mkVarDecl (Variable _ _ (SArray t)) = \i -> do extDeclare $ arrayStruct t - -- return $ arrayDeclaration t i - -- mkVarDecl (Variable _ _ d@(SData _ _)) = \i -> do extDeclare $ datumStruct d - -- return $ datumDeclaration d i - -- mkVarDecl v = error $ "flattenSCon.Lam_.mkVarDecl cannot handle vars of type " ++ show v - - -flattenSCon (PrimOp_ op) = flattenPrimOp op - -flattenSCon (ArrayOp_ op) = flattenArrayOp op - -flattenSCon (Summate _ sr) = - \(lo :* hi :* body :* End) -> - \loc -> - caseBind body $ \v body' -> - do loId <- genIdent - hiId <- genIdent - declare (typeOf lo) loId - declare (typeOf hi) hiId - let loE = CVar loId - hiE = CVar hiId - flattenABT lo loE - flattenABT hi hiE - - iterI <- createIdent v - declare SNat iterI - - accI <- genIdent' "acc" - let semiT = sing_HSemiring sr - declare semiT accI - assign accI (case semiT of - SProb -> negInfinityE - SReal -> floatE 0 - _ -> intE 0) - - let accVar = CVar accI - iterVar = CVar iterI - - - putStat $ opComment "Summate" - -- logSumExp for probabilities - reductionCG CAddOp - accI - (iterVar .=. loE) - (iterVar .<. hiE) - (CUnary CPostIncOp iterVar) $ - do tmpId <- genIdent - declare (typeOf body') tmpId - let tmpE = CVar tmpId - flattenABT body' tmpE - case semiT of - SProb -> logSumExpCG (S.fromList [accVar,tmpE]) accVar - _ -> putStat . CExpr . Just $ (accVar .+=. tmpE) - - putExprStat (loc .=. accVar) - - -flattenSCon (Product _ sr) = - \(lo :* hi :* body :* End) -> - \loc -> - caseBind body $ \v body' -> - do loId <- genIdent - hiId <- genIdent - declare (typeOf lo) loId - declare (typeOf hi) hiId - let loE = CVar loId - hiE = CVar hiId - flattenABT lo loE - flattenABT hi hiE - - iterI <- createIdent v - declare SNat iterI - - accI <- genIdent' "acc" - let semiT = sing_HSemiring sr - declare semiT accI - assign accI (case semiT of - SProb -> floatE 0 - SReal -> floatE 1 - _ -> intE 1) - - let accVar = CVar accI - iterVar = CVar iterI - - putStat $ opComment "Product" - reductionCG (case semiT of - SProb -> CAddOp - _ -> CMulOp) - accI - (iterVar .=. loE) - (iterVar .<. hiE) - (CUnary CPostIncOp iterVar) $ - do tmpId <- genIdent - declare (typeOf body') tmpId - let tmpE = CVar tmpId - flattenABT body' tmpE - putExprStat $ case semiT of - SProb -> CAssign CAddAssOp accVar - _ -> CAssign CMulAssOp accVar - $ tmpE - - putExprStat (loc .=. accVar) - - --------------------- --- SCon Coersions -- --------------------- - --- at this point, only nonrecusive coersions are implemented -flattenSCon (CoerceTo_ ctyp) = - \(e :* End) -> - \loc -> - do eId <- genIdent - let eT = typeOf e - eE = CVar eId - declare eT eId - flattenABT e eE - putExprStat . (CAssign CAssignOp loc) =<< coerceToType ctyp eT eE - where coerceToType - :: Coercion a b - -> Sing (c :: Hakaru) - -> CExpr - -> CodeGen CExpr - coerceToType (CCons p rest) typ = - \e -> primitiveCoerce p typ e >>= coerceToType rest typ - coerceToType CNil _ = return . id - - primitiveCoerce - :: PrimCoercion a b - -> Sing (c :: Hakaru) - -> CExpr - -> CodeGen CExpr - primitiveCoerce (Signed HRing_Int) SNat = nat2int - primitiveCoerce (Signed HRing_Real) SProb = prob2real - primitiveCoerce (Continuous HContinuous_Prob) SNat = nat2prob - primitiveCoerce (Continuous HContinuous_Real) SInt = int2real - primitiveCoerce (Continuous HContinuous_Real) SNat = int2real - primitiveCoerce a b = error $ "flattenSCon CoerceTo_: cannot preform coersion " - ++ show a - ++ " to " - ++ show b - - - -- implementing ONLY functions found in Hakaru.Syntax.AST - nat2int,nat2prob,prob2real,int2real - :: CExpr -> CodeGen CExpr - nat2int = return - nat2prob = \n -> do ident <- genIdent' "p" - declare SProb ident - assign ident . log1p $ n .-. (intE 1) - return (CVar ident) - prob2real = \p -> do ident <- genIdent' "r" - declare SReal ident - assign ident $ (expm1 p) .+. (intE 1) - return (CVar ident) - int2real = return . CCast doubleDecl - - ------------------------------------ --- SCons in the Stochastic Monad -- ------------------------------------ - -flattenSCon (MeasureOp_ op) = flattenMeasureOp op - -flattenSCon Dirac = - \(e :* End) -> - \loc -> - do sId <- genIdent' "samp" - declare (typeOf e) sId - let sE = CVar sId - flattenABT e sE - putExprStat $ mdataPtrWeight loc .=. (floatE 0) - putExprStat $ mdataPtrSample loc .=. sE - -flattenSCon MBind = - \(ma :* b :* End) -> - \loc -> - caseBind b $ \v@(Variable _ _ typ) mb -> - do -- first - mId <- genIdent' "m" - declare (typeOf ma) mId - let mE = CVar mId - flattenABT ma (address mE) - - -- assign that sample to var - vId <- createIdent v - declare typ vId - assign vId (mdataSample mE) - flattenABT mb loc - putExprStat $ mdataPtrWeight loc .+=. (mdataWeight mE) - --- for now plats make use of a global sample -flattenSCon Plate = - \(size :* b :* End) -> - \loc -> - caseBind b $ \v@(Variable _ _ typ) body -> - do sizeId <- genIdent' "s" - declare SNat sizeId - let sizeE = CVar sizeId - flattenABT size sizeE - putExprStat $ (arrayPtrData . mdataPtrSample $ loc) - .=. (CCast (mkPtrDecl . buildType $ typ) - (mkUnary "malloc" - (sizeE .*. (CSizeOfType . mkDecl . buildType $ typ)))) - - - - weightId <- genIdent' "w" - declare SProb weightId - let weightE = CVar weightId - assign weightId (floatE 0) - - itId <- createIdent v - declare SNat itId - let itE = CVar itId - currInd = indirect $ (CMember (mdataSample loc) (Ident "data") True) .+. itE - - sampId <- genIdent' "samp" - declare (typeOf $ body) sampId - let sampE = CVar sampId - - reductionCG CAddOp - weightId - (itE .=. (intE 0)) - (itE .<. sizeE) - (CUnary CPostIncOp itE) - (do flattenABT body (address sampE) - putExprStat (currInd .=. (mdataSample sampE)) - putExprStat (weightE .+=. (mdataWeight sampE))) - - putExprStat $ mdataPtrWeight loc .=. weightE - - ------------------------------------ --- SCon's that arent implemented -- ------------------------------------ - -flattenSCon x = \_ -> \_ -> error $ "TODO: flattenSCon: " ++ show x - - - - - --------------------------------------------------------------------------------- --- NaryOps -- --------------------------------------------------------------------------------- - -flattenNAryOp :: ABT Term abt - => NaryOp a - -> S.Seq (abt '[] a) - -> (CExpr -> CodeGen ()) -flattenNAryOp op args = - \loc -> - do es <- T.forM args $ \a -> - do aId <- genIdent - let aE = CVar aId - declare (typeOf a) aId - _ <- flattenABT a aE - return aE - case op of - And -> boolNaryOp op es loc - Or -> boolNaryOp op es loc - Xor -> boolNaryOp op es loc - Iff -> boolNaryOp op es loc - - (Sum HSemiring_Prob) -> logSumExpCG es loc - - _ -> let opE = F.foldr (binaryOp op) (S.index es 0) (S.drop 1 es) - in putExprStat (loc .=. opE) - - - where boolNaryOp op' es' loc' = - let indexOf x = CMember x (Ident "index") True - es'' = fmap indexOf es' - expr = F.foldr (binaryOp op') - (S.index es'' 0) - (S.drop 1 es'') - in putExprStat ((indexOf loc') .=. expr) - - --------------------------------------- --- LogSumExp for NaryOp Add [SProb] -- --------------------------------------- -{- - -Special for addition of probabilities we have a logSumExp. This will compute the -sum of the probabilities safely. Just adding the exp(a . prob) would make us -loose any of the safety from underflow that we got from storing prob in the log -domain - --} - --- the tree traversal is a depth first search -logSumExp :: S.Seq CExpr -> CExpr -logSumExp es = mkCompTree 0 1 - - where lastIndex = S.length es - 1 - - compIndices :: Int -> Int -> CExpr -> CExpr -> CExpr - compIndices i j = CCond ((S.index es i) .>. (S.index es j)) - - mkCompTree :: Int -> Int -> CExpr - mkCompTree i j - | j == lastIndex = compIndices i j (logSumExp' i) (logSumExp' j) - | otherwise = compIndices i j - (mkCompTree i (succ j)) - (mkCompTree j (succ j)) - - diffExp :: Int -> Int -> CExpr - diffExp a b = expm1 ((S.index es a) .-. (S.index es b)) - - -- given the max index, produce a logSumExp expression - logSumExp' :: Int -> CExpr - logSumExp' 0 = S.index es 0 - .+. (log1p $ foldr (\x acc -> diffExp x 0 .+. acc) - (diffExp 1 0) - [2..S.length es - 1] - .+. (intE $ fromIntegral lastIndex)) - logSumExp' i = S.index es i - .+. (log1p $ foldr (\x acc -> if i == x - then acc - else diffExp x i .+. acc) - (diffExp 0 i) - [1..S.length es - 1] - .+. (intE $ fromIntegral lastIndex)) - - --- | logSumExpCG creates global functions for every n-ary logSumExp function --- this reduces code size -logSumExpCG :: S.Seq CExpr -> (CExpr -> CodeGen ()) -logSumExpCG seqE = - let size = S.length $ seqE - name = "logSumExp" ++ (show size) - funcId = Ident name - in \loc -> do-- reset the names so that the function is the same for each arity - cg <- get - put (cg { freshNames = suffixes }) - argIds <- replicateM size genIdent - let decls = fmap (typeDeclaration SProb) argIds - vars = fmap CVar argIds - extDeclare . CFunDefExt $ functionDef SProb - funcId - decls - [] - [CReturn . Just $ logSumExp $ S.fromList vars ] - cg' <- get - put (cg' { freshNames = freshNames cg }) - putExprStat $ loc .=. (CCall (CVar funcId) (F.toList seqE)) - - --------------------------------------------------------------------------------- --- Literals -- --------------------------------------------------------------------------------- - -flattenLit - :: Literal a - -> (CExpr -> CodeGen ()) -flattenLit lit = - \loc -> - case lit of - (LNat x) -> putExprStat $ loc .=. (intE $ fromIntegral x) - (LInt x) -> putExprStat $ loc .=. (intE x) - (LReal x) -> putExprStat $ loc .=. (floatE $ fromRational x) - (LProb x) -> let rat = fromNonNegativeRational x - x' = (fromIntegral $ numerator rat) - / (fromIntegral $ denominator rat) - xE = log1p (floatE x' .-. intE 1) - in putExprStat (loc .=. xE) - --------------------------------------------------------------------------------- --- Array and ArrayOps -- --------------------------------------------------------------------------------- - - -flattenArray - :: (ABT Term abt) - => (abt '[] 'HNat) - -> (abt '[ 'HNat ] a) - -> (CExpr -> CodeGen ()) -flattenArray arity body = - \loc -> - caseBind body $ \v@(Variable _ _ typ) body' -> - let arityE = arraySize loc - dataE = arrayData loc in - do flattenABT arity arityE - putExprStat $ dataE - .=. (CCast (mkPtrDecl . buildType $ typ) - (mkUnary "malloc" - (arityE .*. (CSizeOfType . mkDecl . buildType $ typ)))) - - itId <- createIdent v - declare SNat itId - let itE = CVar itId - currInd = indirect (dataE .+. itE) - - putStat $ opComment "Create Array" - forCG (itE .=. (intE 0)) - (itE .<. arityE) - (CUnary CPostIncOp itE) - (flattenABT body' currInd) - - --------------- --- ArrayOps -- --------------- - -flattenArrayOp - :: ( ABT Term abt - , typs ~ UnLCs args - , args ~ LCs typs - ) - => ArrayOp typs a - -> SArgs abt args - -> (CExpr -> CodeGen ()) - - -flattenArrayOp (Index _) = - \(arr :* ind :* End) -> - \loc -> - do arrId <- genIdent' "arr" - indId <- genIdent - let arrE = CVar arrId - indE = CVar indId - declare (typeOf arr) arrId - declare SNat indId - flattenABT arr arrE - flattenABT ind indE - let valE = indirect ((CMember arrE (Ident "data") True) .+. indE) - putExprStat (loc .=. valE) - -flattenArrayOp (Size _) = - \(arr :* End) -> - \loc -> - do arrId <- genIdent' "arr" - declare (typeOf arr) arrId - let arrE = CVar arrId - flattenABT arr arrE - putExprStat (loc .=. (CMember arrE (Ident "size") True)) - -flattenArrayOp (Reduce _) = error "TODO: flattenArrayOp" - -- \(fun :* base :* arr :* End) -> - -- do funE <- flattenABT fun - -- baseE <- flattenABT base - -- arrE <- flattenABT arr - -- accI <- genIdent' "acc" - -- iterI <- genIdent' "iter" - - -- let sizeE = CMember arrE (Ident "size") True - -- iterE = CVar iterI - -- accE = CVar accI - -- cond = iterE .<. sizeE - -- inc = CUnary CPostIncOp iterE - - -- declare (typeOf base) accI - -- declare SInt iterI - -- assign accI baseE - -- forCG (iterE .=. (intE 0)) cond inc $ - -- assign accI $ CCall funE [accE] - - -- return accE - - --------------------------------------------------------------------------------- --- Datum and Case -- --------------------------------------------------------------------------------- -{- - -Datum are sums of products of types. This maps to a C structure. flattenDatum -will produce a literal of some datum type. This will also produce a global -struct representing that datum which will be needed for the C compiler. - --} - - -flattenDatum - :: (ABT Term abt) - => Datum (abt '[]) (HData' a) - -> (CExpr -> CodeGen ()) -flattenDatum (Datum _ typ code) = - \loc -> - do extDeclare $ datumStruct typ - assignDatum code loc - -datumNames :: [String] -datumNames = filter (\n -> not $ elem (head n) ['0'..'9']) names - where base = ['0'..'9'] ++ ['a'..'z'] - names = [[x] | x <- base] `mplus` (do n <- names - [n++[x] | x <- base]) - -assignDatum - :: (ABT Term abt) - => DatumCode xss (abt '[]) c - -> CExpr - -> CodeGen () -assignDatum code ident = - let index = getIndex code - indexExpr = CMember ident (Ident "index") True - in do putExprStat (indexExpr .=. (intE index)) - sequence_ $ assignSum code ident - where getIndex :: DatumCode xss b c -> Integer - getIndex (Inl _) = 0 - getIndex (Inr rest) = succ (getIndex rest) - -assignSum - :: (ABT Term abt) - => DatumCode xs (abt '[]) c - -> CExpr - -> [CodeGen ()] -assignSum code ident = fst $ runState (assignSum' code ident) datumNames - -assignSum' - :: (ABT Term abt) - => DatumCode xs (abt '[]) c - -> CExpr - -> State [String] [CodeGen ()] -assignSum' (Inr rest) topIdent = - do (_:names) <- get - put names - assignSum' rest topIdent -assignSum' (Inl prod) topIdent = - do (name:_) <- get - return $ assignProd prod topIdent (CVar . Ident $ name) - -assignProd - :: (ABT Term abt) - => DatumStruct xs (abt '[]) c - -> CExpr - -> CExpr - -> [CodeGen ()] -assignProd dstruct topIdent sumIdent = - fst $ runState (assignProd' dstruct topIdent sumIdent) datumNames - -assignProd' - :: (ABT Term abt) - => DatumStruct xs (abt '[]) c - -> CExpr - -> CExpr - -> State [String] [CodeGen ()] -assignProd' Done _ _ = return [] -assignProd' (Et (Konst d) rest) topIdent (CVar sumIdent) = - do (name:names) <- get - put names - let varName = CMember (CMember (CMember topIdent - (Ident "sum") - True) - sumIdent - True) - (Ident name) - True - rest' <- assignProd' rest topIdent (CVar sumIdent) - return $ [flattenABT d varName] ++ rest' -assignProd' _ _ _ = error $ "TODO: assignProd Ident" - - ----------- --- Case -- ----------- - --- currently we can only match on boolean values -flattenCase - :: forall abt a b - . (ABT Term abt) - => abt '[] a - -> [Branch a abt b] - -> (CExpr -> CodeGen ()) - -flattenCase c (Branch (PDatum _ (PInl PDone)) trueB:Branch (PDatum _ (PInr (PInl PDone))) falseB:[]) = - \loc -> - do cId <- genIdent - declare (typeOf c) cId - let cE = (CVar cId) - flattenABT c cE - - cg <- get - let trueM = flattenABT trueB loc - falseM = flattenABT falseB loc - (_,cg') = runState trueM $ cg { statements = [] } - (_,cg'') = runState falseM $ cg' { statements = [] } - put $ cg'' { statements = statements cg } - - putStat $ CIf ((CMember cE (Ident "index") True) .==. (intE 0)) - (CCompound . fmap CBlockStat . reverse . statements $ cg') - Nothing - putStat $ CIf ((CMember cE (Ident "index") True) .==. (intE 1)) - (CCompound . fmap CBlockStat . reverse . statements $ cg'') - Nothing - - -flattenCase _ _ = error "TODO: flattenCase" - - --------------------------------------------------------------------------------- --- PrimOp -- --------------------------------------------------------------------------------- - -flattenPrimOp - :: ( ABT Term abt - , typs ~ UnLCs args - , args ~ LCs typs) - => PrimOp typs a - -> SArgs abt args - -> (CExpr -> CodeGen ()) - - -flattenPrimOp Pi = - \End -> - \loc -> let piE = log1p ((CVar . Ident $ "M_PI") .-. (intE 1)) in - putExprStat (loc .=. piE) - -flattenPrimOp Not = - \(a :* End) -> - \_ -> - -- this is currently incorrect, need to use memcpy to preserve value of - -- 'a' - do tmpId <- genIdent' "not" - declare sBool tmpId - let tmpE = CVar tmpId - flattenABT a tmpE - let datumIndex = CMember tmpE (Ident "index") True - putExprStat $ datumIndex .=. (CCond (datumIndex .==. (intE 1)) - (intE 0) - (intE 1)) - -flattenPrimOp RealPow = - \(base :* power :* End) -> - \loc -> - do baseId <- genIdent - powerId <- genIdent - declare SProb baseId - declare SReal powerId - let baseE = CVar baseId - powerE = CVar powerId - flattenABT base baseE -- first argument is a Prob - flattenABT power powerE - let realPow = CCall (CVar . Ident $ "pow") - [ expm1 baseE .+. (intE 1), powerE] - putExprStat $ loc .=. (log1p (realPow .-. (intE 1))) - -flattenPrimOp (NatPow baseTyp) = - \(base :* power :* End) -> - \loc -> - let sBase = sing_HSemiring baseTyp in - do baseId <- genIdent - powerId <- genIdent - declare sBase baseId - declare SReal powerId - let baseE = CVar baseId - powerE = CVar powerId - flattenABT base baseE - flattenABT power powerE - let powerOf x y = CCall (CVar . Ident $ "pow") [x,y] - value = case sBase of - SProb -> log1p $ (powerOf (expm1 baseE .+. (intE 1)) powerE) - .-. (intE 1) - _ -> powerOf baseE powerE - putExprStat $ loc .=. value - -flattenPrimOp (NatRoot baseTyp) = - \(base :* root :* End) -> - \loc -> - let sBase = sing_HRadical baseTyp in - do baseId <- genIdent - rootId <- genIdent - declare sBase baseId - declare SReal rootId - let baseE = CVar baseId - rootE = CVar rootId - flattenABT base baseE - flattenABT root rootE - let powerOf x y = CCall (CVar . Ident $ "pow") [x,y] - recipE = (floatE 1) ./. rootE - value = case sBase of - SProb -> log1p $ (powerOf (expm1 baseE .+. (intE 1)) recipE) - .-. (intE 1) - _ -> powerOf baseE recipE - putExprStat $ loc .=. value - -flattenPrimOp (Recip t) = - \(a :* End) -> - \loc -> - do aId <- genIdent - declare (typeOf a) aId - let aE = CVar aId - flattenABT a aE - case t of - HFractional_Real -> putExprStat $ loc .=. ((intE 1) ./. aE) - HFractional_Prob -> putExprStat $ loc .=. (CUnary CMinOp aE) - --- | exp : real -> prob, because of this we can just turn it into a prob without taking --- its log, which would give us an exp in the log-domain -flattenPrimOp Exp = \(a :* End) -> flattenABT a - -flattenPrimOp (Equal _) = - \(a :* b :* End) -> - \loc -> - do aId <- genIdent - bId <- genIdent - let aE = CVar aId - bE = CVar bId - aT = typeOf a - bT = typeOf b - declare aT aId - declare bT bId - flattenABT a aE - flattenABT b bE - - -- special case for booleans - let aE' = case aT of - (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CMember aE (Ident "index") True) - _ -> aE - let bE' = case bT of - (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CMember bE (Ident "index") True) - _ -> bE - - putExprStat $ (CMember loc (Ident "index") True) - .=. (CCond (aE' .==. bE') (intE 0) (intE 1)) - - -flattenPrimOp (Less _) = - \(a :* b :* End) -> - \loc -> - do aId <- genIdent - bId <- genIdent - let aE = CVar aId - bE = CVar bId - declare (typeOf a) aId - declare (typeOf b) bId - flattenABT a aE - flattenABT b bE - putExprStat $ (CMember loc (Ident "index") True) - .=. (CCond (aE .<. bE) (intE 0) (intE 1)) - -flattenPrimOp (Negate HRing_Real) = - \(a :* End) -> - \loc -> - do negId <- genIdent' "neg" - declare SReal negId - let negE = CVar negId - flattenABT a negE - putExprStat $ loc .=. (CUnary CMinOp $ negE) - - -flattenPrimOp t = \_ -> error $ "TODO: flattenPrimOp: " ++ show t - - --------------------------------------------------------------------------------- --- MeasureOps and Superpose -- --------------------------------------------------------------------------------- - -{- - -The sections contains operations in the stochastic monad. See also -(Dirac, MBind, and Plate) found in SCon. Also see Reject found at the top level. - -Remember in the C runtime. Measures are housed in a measure function, which -takes an `struct mdata` location. The MeasureOp attempts to store a value at -that location and returns 0 if it fails and 1 if it succeeds in that task. - -The functions uniformCG, normalCG, and gammaCG are primitives that will generate -functions and call them (similar to logSumExpCG). The reduce code size and make -samplers a little more readable. - -TODO: add inline pragmas to uniformCG, normalCG, and gammaCG --} - -uniformFun :: CFunDef -uniformFun = CFunDef [CTypeSpec CVoid] - (CDeclr Nothing [CDDeclrIdent funcId]) - [typeDeclaration SReal loId - ,typeDeclaration SReal hiId - ,typePtrDeclaration (SMeasure SReal) mId] - (seqCStat $ comment ++[assW,assS,CReturn Nothing]) - where r = CCast doubleDecl rand - rMax = CCast doubleDecl (CVar . Ident $ "RAND_MAX") - (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) - (loId,loE) = let ident = Ident "lo" in (ident,CVar ident) - (hiId,hiE) = let ident = Ident "hi" in (ident,CVar ident) - value = (loE .+. ((r ./. rMax) .*. (hiE .-. loE))) - comment = fmap CComment - ["uniform :: real -> real -> *(mdata real) -> ()" - ,"------------------------------------------------"] - assW = CExpr . Just $ mdataPtrWeight mE .=. (floatE 0) - assS = CExpr . Just $ mdataPtrSample mE .=. value - funcId = Ident "uniform" - - -uniformCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) -uniformCG aE bE = - \loc -> do - reserveName "uniform" - extDeclare . CFunDefExt $ uniformFun - putExprStat $ CCall (CVar . Ident $ "uniform") [aE,bE,loc] - - -{- - This is very cryptic, but I assure you it is only building an AST for the - Marsaglia Polar Method --} - -normalFun :: CFunDef -normalFun = CFunDef [CTypeSpec CVoid] - (CDeclr Nothing [CDDeclrIdent (Ident "normal")]) - [typeDeclaration SReal aId - ,typeDeclaration SProb bId - ,typePtrDeclaration (SMeasure SReal) mId] - (CCompound $ comment ++ decls ++ stmts) - - where r = CCast doubleDecl rand - rMax = CCast doubleDecl (CVar . Ident $ "RAND_MAX") - (aId,aE) = let ident = Ident "a" in (ident,CVar ident) - (bId,bE) = let ident = Ident "b" in (ident,CVar ident) - (qId,qE) = let ident = Ident "q" in (ident,CVar ident) - (uId,uE) = let ident = Ident "u" in (ident,CVar ident) - (vId,vE) = let ident = Ident "v" in (ident,CVar ident) - (rId,rE) = let ident = Ident "r" in (ident,CVar ident) - (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) - draw xE = CExpr . Just $ xE .=. (((r ./. rMax) .*. (floatE 2)) .-. (floatE 1)) - body = seqCStat [draw uE - ,draw vE - ,CExpr . Just $ qE .=. ((uE .*. uE) .+. (vE .*. vE))] - polar = CWhile (qE .>. (floatE 1)) body True - setR = CExpr . Just $ rE .=. (sqrt (((CUnary CMinOp (floatE 2)) .*. log qE) ./. qE)) - finalValue = aE .+. (uE .*. rE .*. bE) - comment = fmap (CBlockStat . CComment) - ["normal :: real -> real -> *(mdata real) -> ()" - ,"Marsaglia Polar Method" - ,"-----------------------------------------------"] - decls = fmap (CBlockDecl . typeDeclaration SReal) [uId,vId,qId,rId] - stmts = fmap CBlockStat [polar,setR, assW, assS,CReturn Nothing] - assW = CExpr . Just $ mdataPtrWeight mE .=. (floatE 0) - assS = CExpr . Just $ mdataPtrSample mE .=. finalValue - - -normalCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) -normalCG aE bE = - \loc -> do - reserveName "normal" - extDeclare . CFunDefExt $ normalFun - putExprStat $ CCall (CVar . Ident $ "normal") [aE,bE,loc] - -{- - This method is from Marsaglia and Tsang "a simple method for generating gamma variables" --} -gammaFun :: CFunDef -gammaFun = CFunDef [CTypeSpec CVoid] - (CDeclr Nothing [CDDeclrIdent (Ident "gamma")]) - [typeDeclaration SProb aId - ,typeDeclaration SProb bId - ,typePtrDeclaration (SMeasure SProb) mId] - (CCompound $ comment ++ decls ++ stmts) - where (aId,aE) = let ident = Ident "a" in (ident,CVar ident) - (bId,bE) = let ident = Ident "b" in (ident,CVar ident) - (cId,cE) = let ident = Ident "c" in (ident,CVar ident) - (dId,dE) = let ident = Ident "d" in (ident,CVar ident) - (xId,xE) = let ident = Ident "x" in (ident,CVar ident) - (vId,vE) = let ident = Ident "v" in (ident,CVar ident) - (uId,uE) = let ident = Ident "u" in (ident,CVar ident) - (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) - comment = fmap (CBlockStat . CComment) - ["gamma :: real -> prob -> *(mdata prob) -> ()" - ,"Marsaglia and Tsang 'a simple method for generating gamma variables'" - ,"--------------------------------------------------------------------"] - decls = fmap CBlockDecl $ (fmap (typeDeclaration SReal) [dId,cId,vId]) - ++ (fmap (typeDeclaration (SMeasure SReal)) [uId,xId]) - stmts = fmap CBlockStat $ [assD,assC,outerWhile] - xS = mdataSample xE - uS = mdataSample uE - assD = CExpr . Just $ dE .=. (aE .-. ((floatE 1) ./. (floatE 3))) - assC = CExpr . Just $ cE .=. ((floatE 1) ./. (sqrt ((floatE 9) .*. dE))) - outerWhile = CWhile (intE 1) (seqCStat [innerWhile,assV,assU,exit]) False - innerWhile = CWhile (vE .<=. (floatE 0)) (seqCStat [assX,assVIn]) True - assX = CExpr . Just $ CCall (CVar . Ident $ "normal") [(floatE 0),(floatE 1),address xE] - assVIn = CExpr . Just $ vE .=. ((floatE 1) .+. (cE .*. xS)) - assV = CExpr . Just $ vE .=. (vE .*. vE .*. vE) - assU = CExpr . Just $ CCall (CVar . Ident $ "uniform") [(floatE 0),(floatE 1),address uE] - exitC1 = uS .<. ((floatE 1) .-. ((floatE 0.331 .*. (xS .*. xS) .*. (xS .*. xS)))) - exitC2 = (log uS) .<. (((floatE 0.5) .*. (xS .*. xS)) .+. (dE .*. ((floatE 1.0) .-. vE .+. (log vE)))) - assW = CExpr . Just $ mdataPtrWeight mE .=. (floatE 0) - assS = CExpr . Just $ mdataPtrSample mE .=. (log (dE .*. vE)) .+. bE - exit = CIf (exitC1 .||. exitC2) (seqCStat [assW,assS,CReturn Nothing]) Nothing - - -gammaCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) -gammaCG aE bE = - \loc -> do - extDeclare $ mdataStruct SReal - mapM_ reserveName ["uniform","normal","gamma"] - mapM_ (extDeclare . CFunDefExt) [uniformFun,normalFun,gammaFun] - putExprStat $ CCall (CVar . Ident $ "gamma") [aE,bE,loc] - - -flattenMeasureOp - :: forall abt typs args a . - ( ABT Term abt - , typs ~ UnLCs args - , args ~ LCs typs ) - => MeasureOp typs a - -> SArgs abt args - -> (CExpr -> CodeGen ()) - - -flattenMeasureOp Uniform = - \(a :* b :* End) -> - \loc -> - do (aId:bId:[]) <- replicateM 2 genIdent - let aE = CVar aId - bE = CVar bId - declare SReal aId - declare SReal bId - flattenABT a aE - flattenABT b bE - uniformCG aE bE loc - - -flattenMeasureOp Normal = - \(a :* b :* End) -> - \loc -> - do (aId:bId:[]) <- replicateM 2 genIdent - let aE = CVar aId - bE = CVar bId - declare SReal aId - declare SReal bId - flattenABT a aE - flattenABT b bE - normalCG aE (exp bE) loc - - -flattenMeasureOp Gamma = - \(a :* b :* End) -> - \loc -> - do (aId:bId:[]) <- replicateM 2 genIdent - let aE = CVar aId - bE = CVar bId - declare SReal aId - declare SReal bId - flattenABT a aE - flattenABT b bE - gammaCG (exp aE) bE loc - - -flattenMeasureOp Beta = - \(a :* b :* End) -> flattenABT (HKP.beta'' a b) - - -flattenMeasureOp Categorical = \(arr :* End) -> - \loc -> - do arrId <- genIdent - declare (typeOf arr) arrId - let arrE = CVar arrId - flattenABT arr arrE - - itId <- genIdent' "it" - declare SInt itId - let itE = CVar itId - - wSumId <- genIdent' "ws" - declare SProb wSumId - let wSumE = CVar wSumId - assign wSumId (log (intE 0)) - - let currE = indirect (arrayData arrE .+. itE) - cond = itE .<. (arraySize arrE) - inc = CUnary CPostIncOp itE - - isPar <- isParallel - mkSequential - - -- first calculate the max weight - forCG (itE .=. (intE 0)) cond inc $ - logSumExpCG (S.fromList [wSumE,currE]) wSumE - - -- draw number from uniform(0, weightSum) - rId <- genIdent' "r" - declare SReal rId - let r = CCast doubleDecl rand - rMax = CCast doubleDecl (CVar . Ident $ "RAND_MAX") - rE = CVar rId - assign rId ((r ./. rMax) .*. (exp wSumE)) - - assign wSumId (log (intE 0)) - assign itId (intE 0) - whileCG (intE 1) - $ do stat <- runCodeGenBlock $ - do putExprStat $ mdataPtrWeight loc .=. (intE 0) - putExprStat $ mdataPtrSample loc .=. itE - putStat CBreak - putStat $ CIf (rE .<. (exp wSumE)) stat Nothing - logSumExpCG (S.fromList [wSumE,currE]) wSumE - putExprStat $ CUnary CPostIncOp itE - - when isPar mkParallel - - -flattenMeasureOp x = error $ "TODO: flattenMeasureOp: " ++ show x - ---------------- --- Superpose -- ---------------- - -flattenSuperpose - :: (ABT Term abt) - => NE.NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a)) - -> (CExpr -> CodeGen ()) - --- do we need to normalize? -flattenSuperpose pairs = - let pairs' = NE.toList pairs in - - if length pairs' == 1 - then \loc -> let (w,m) = head pairs' in - do mId <- genIdent - wId <- genIdent - declare (typeOf m) mId - declare SProb wId - let mE = address . CVar $ mId - wE = CVar wId - flattenABT w wE - flattenABT m mE - putExprStat $ mdataPtrWeight loc .=. ((mdataPtrWeight mE) .+. wE) - putExprStat $ mdataPtrSample loc .=. (mdataPtrSample mE) - - else \loc -> - do wEs <- forM pairs' $ \(w,_) -> - do wId <- genIdent' "w" - declare SProb wId - let wE = CVar wId - flattenABT w wE - return wE - - wSumId <- genIdent' "ws" - declare SProb wSumId - let wSumE = CVar wSumId - logSumExpCG (S.fromList wEs) wSumE - - -- draw number from uniform(0, weightSum) - rId <- genIdent' "r" - declare SReal rId - let r = CCast doubleDecl rand - rMax = CCast doubleDecl (CVar . Ident $ "RAND_MAX") - rE = CVar rId - assign rId ((r ./. rMax) .*. (exp wSumE)) - - -- an iterator for picking a measure - itId <- genIdent' "it" - declare SProb itId - let itE = CVar itId - assign itId (log (intE 0)) - - -- an output measure to assign to - outId <- genIdent' "out" - declare (typeOf . snd . head $ pairs') outId - let outE = address $ CVar outId - - outLabel <- genIdent' "exit" - - forM_ (zip wEs pairs') - $ \(wE,(_,m)) -> - do logSumExpCG (S.fromList [itE,wE]) itE - stat <- runCodeGenBlock (flattenABT m outE >> putStat (CGoto outLabel)) - putStat $ CIf (rE .<. (exp itE)) stat Nothing - - putStat $ CLabel outLabel (CExpr Nothing) - putExprStat $ mdataPtrWeight loc .=. ((mdataPtrWeight outE) .+. wSumE) - putExprStat $ mdataPtrSample loc .=. (mdataPtrSample outE) + , flattenTerm + , flattenWithName + , flattenWithName' + , localVar + , localVar' + , opComment + ) where + +import Language.Hakaru.CodeGen.CodeGenMonad +import Language.Hakaru.CodeGen.AST +import Language.Hakaru.CodeGen.Libs +import Language.Hakaru.CodeGen.Types + +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.TypeOf +import Language.Hakaru.Syntax.Datum hiding (Ident) +import Language.Hakaru.Syntax.Reducer +import qualified Language.Hakaru.Syntax.Prelude as HKP +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.HClasses +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Types.Coercion +import Language.Hakaru.Types.Sing + +import Control.Monad.State.Strict +import Control.Monad (replicateM) +import Control.Applicative (pure) +import Data.Number.Natural +import Data.Monoid hiding (Product,Sum) +import Data.Ratio +import qualified Data.List.NonEmpty as NE +import qualified Data.Sequence as S +import qualified Data.Foldable as F +import qualified Data.Traversable as T + + +#if __GLASGOW_HASKELL__ < 710 +import Data.Functor +#endif + + +opComment :: String -> CStat +opComment opStr = CComment $ concat [space," ",opStr," ",space] + where size = (50 - (length opStr)) `div` 2 - 8 + space = replicate size '-' + +-------------------------------------------------------------------------------- +-- Top Level -- +-------------------------------------------------------------------------------- +{- + +flattening an ABT will produce a continuation that takes a CExpr representing +a location where the value of the ABT should be stored. Return type of the +the continuation is CodeGen Bool, where the computed bool is whether or not +there is a Reject inside the ABT. Therefore it is only needed when computing +mochastic values + +-} + +localVar :: Sing (a :: Hakaru) -> CodeGen CExpr +localVar typ = localVar' typ "" + +localVar' :: Sing (a :: Hakaru) -> String -> CodeGen CExpr +localVar' typ s = + do eId <- genIdent' s + declare typ eId + return (CVar eId) + + +flattenWithName' + :: ABT Term abt + => abt '[] a + -> String + -> CodeGen CExpr +flattenWithName' abt hint = do + ident <- genIdent' hint + declare (typeOf abt) ident + let cvar = CVar ident + flattenABT abt cvar + return cvar + +flattenWithName + :: ABT Term abt + => abt '[] a + -> CodeGen CExpr +flattenWithName abt = flattenWithName' abt "" + +flattenABT + :: ABT Term abt + => abt '[] a + -> (CExpr -> CodeGen ()) +flattenABT abt = caseVarSyn abt flattenVar flattenTerm + +-- note that variables will find their values in the state of the CodeGen monad +flattenVar + :: Variable (a :: Hakaru) + -> (CExpr -> CodeGen ()) +flattenVar v = \loc -> + do v' <- CVar <$> lookupIdent v + putExprStat $ loc .=. v' + +flattenTerm + :: ABT Term abt + => Term abt a + -> (CExpr -> CodeGen ()) +-- SCon can contain mochastic terms +flattenTerm (x :$ ys) = flattenSCon x ys + +flattenTerm (NaryOp_ t s) = flattenNAryOp t s +flattenTerm (Literal_ x) = flattenLit x +flattenTerm (Empty_ _) = error "TODO: flattenTerm{Empty}" + +flattenTerm (Datum_ d) = flattenDatum d +flattenTerm (Case_ c bs) = flattenCase c bs + +flattenTerm (Bucket b e rs) = flattenBucket b e rs + +flattenTerm (Array_ s e) = flattenArray s e +flattenTerm (ArrayLiteral_ s) = flattenArrayLiteral s + + +--------------------- +-- Mochastic Terms -- +--------------------- +flattenTerm (Superpose_ wes) = flattenSuperpose wes +flattenTerm (Reject_ _) = \loc -> putExprStat (mdataWeight loc .=. (intE 0)) + + +-------------------------------------------------------------------------------- +-- SCon -- +-------------------------------------------------------------------------------- + +flattenSCon + :: ( ABT Term abt ) + => SCon args a + -> SArgs abt args + -> (CExpr -> CodeGen ()) + +flattenSCon Let_ = + \(expr :* body :* End) -> + \loc -> do + caseBind body $ \v@(Variable _ _ typ) body'-> + do ident <- createIdent v + case typ of + (SFun _ _) -> return () + _ -> declare typ ident + flattenABT expr (CVar ident) + flattenABT body' loc + +-- Lambdas produce functions and then return a function label exprssion +flattenSCon Lam_ = + \(body :* End) -> + \loc -> do + -- externally declare closure and function + closureTypeSpec <- coalesceLambda body extDeclClosure + + -- declare local closure var + closureId <- genIdent' "closure" + declare' (buildDeclaration closureTypeSpec closureId) + + -- capture environment in closure + putExprStat $ loc .=. (CVar closureId) + + where coalesceLambda + :: ( ABT Term abt ) + => abt '[x] a + -> (forall (ys :: [Hakaru]) b. List1 Variable ys -> abt '[] b -> r) + -> r + coalesceLambda abt k = + caseBind abt $ \v abt' -> + caseVarSyn abt' (const (k (Cons1 v Nil1) abt')) $ \term -> + case term of + (Lam_ :$ body :* End) -> + coalesceLambda body $ \vars abt'' -> k (Cons1 v vars) abt'' + _ -> k (Cons1 v Nil1) abt' + + -- given a parameter, create identifiers corresponding to Hakaru vars, + -- and return a CTypeSpec for the param + -- Will this fail if the parameter is a SFun? + mkVarIdandSpec :: Variable (a :: Hakaru) -> CodeGen (Ident,[CTypeSpec]) + mkVarIdandSpec v@(Variable _ _ typ) = do + extDeclareTypes typ + vId <- createIdent v + return (vId,buildType typ) + + extDeclClosure + :: ( ABT Term abt ) + => List1 Variable (ys :: [Hakaru]) + -> abt '[] b + -> CodeGen CTypeSpec + extDeclClosure vars body'= do + funcId <- genIdent' "fn" + idAndSpecs <- sequence $ foldMap11 (\v -> [mkVarIdandSpec v]) vars + let fVars = freeVars body' + typ = typeOf body' + sId@(Ident sname) <- extDeclClosureStruct typ (fmap snd idAndSpecs) fVars + funCG (head . buildType $ typ) + funcId + ([buildDeclaration (callStruct sname) (Ident "env")] + ++ (fmap (\(vId,specs) -> buildDeclaration' specs vId) idAndSpecs)) + ((putStat . CReturn . Just) =<< flattenWithName body') + return (callStruct sname) + + extDeclClosureStruct + :: forall (a :: Hakaru) (ys :: [Hakaru]) + . Sing a + -> [[CTypeSpec]] + -> VarSet (KindOf a) + -> CodeGen Ident + extDeclClosureStruct retTyp paramTypeSpecs freeVars = do + sId@(Ident sname) <- genIdent' "clos" + freeVarDecls <- mapM (\(SomeVariable v@(Variable _ _ typ)) -> do + extDeclareTypes typ + vId <- createIdent v + return (typeDeclaration typ vId) + ) (fromVarSet freeVars) + let funPtrDecl = + CDecl (fmap CTypeSpec $ buildType retTyp) + [( CDeclr Nothing + (CDDeclrFun + (CDDeclrRec (CDeclr (Just $ CPtrDeclr []) (CDDeclrIdent . Ident $ "fn"))) + ([callStruct sname]++(concat paramTypeSpecs))) + , Nothing)] + extDeclare $ CDeclExt + $ CDecl [ CTypeSpec $ buildStruct (Just sId) + ([funPtrDecl]++freeVarDecls) ] + [] + return sId + +flattenSCon App_ = + \(fun :* arg :* End) -> + \loc -> do + closE <- flattenWithName' fun "clos" + paramE <- flattenWithName' fun "param" + putExprStat $ loc .=. (CCall (CMember closE (Ident "fn") True) [paramE]) + +flattenSCon (PrimOp_ op) = flattenPrimOp op + +flattenSCon (ArrayOp_ op) = flattenArrayOp op + +flattenSCon (Summate _ sr) = + \(lo :* hi :* body :* End) -> + \loc -> + let semiTyp = sing_HSemiring sr in + do loE <- flattenWithName' lo "lo" + hiE <- flattenWithName' hi "hi" + + putStat $ opComment "Begin Summate" + + case semiTyp of + -- special prob branch + SProb -> do + summateArrId <- genIdent' "summate_arr" + declare (SArray SProb) summateArrId + let summateArrE = CVar summateArrId + putExprStat $ arraySize summateArrE .=. (hiE .-. loE) + putExprStat $ arrayData summateArrE + .=. (castToPtrOf [CDouble] + (mallocE ((arraySize summateArrE) .*. + (CSizeOfType (CTypeName [CDouble] False))))) + lseSummateArrayCG body summateArrE loc + putExprStat $ freeE (arrayData summateArrE) + + _ -> + caseBind body $ \v body' -> do + iterI <- createIdent v + declare SNat iterI + + accI <- genIdent' "acc" + declare semiTyp accI + assign accI (case semiTyp of + SReal -> floatE 0 + _ -> intE 0) + + let accVar = CVar accI + iterVar = CVar iterI + + reductionCG (Left CAddOp) + accI + (iterVar .=. loE) + (iterVar .<. hiE) + (CUnary CPostIncOp iterVar) $ + (putExprStat . (accVar .+=.) =<< flattenWithName body') + putExprStat $ loc .=. accVar + + putStat $ opComment "End Summate" + + + +flattenSCon (Product _ sr) = + \(lo :* hi :* body :* End) -> + \loc -> + let semiTyp = sing_HSemiring sr in + do loE <- flattenWithName' lo "lo" + hiE <- flattenWithName' hi "hi" + + putStat $ opComment "Begin Product" + + case semiTyp of + -- special prob branch + SProb -> kahanSummationCG body loE hiE loc + + _ -> + caseBind body $ \v body' -> do + iterI <- createIdent v + declare SNat iterI + + accI <- genIdent' "acc" + declare semiTyp accI + assign accI (case semiTyp of + SReal -> floatE 1 + _ -> intE 1) + + let accVar = CVar accI + iterVar = CVar iterI + + reductionCG (Left CMulOp) + accI + (iterVar .=. loE) + (iterVar .<. hiE) + (CUnary CPostIncOp iterVar) $ + (putExprStat . (accVar .*=.) =<< flattenWithName body') + putExprStat (loc .=. accVar) + putStat $ opComment "End Product" + + +-------------------- +-- SCon Coercions -- +-------------------- +{- Helpers found by searching "Coercion Helpers" -} + +flattenSCon (CoerceTo_ ctyp) = + \(e :* End) -> + \loc -> + do eE <- flattenWithName e + cE <- coerceToCG ctyp eE + putExprStat $ loc .=. cE + +flattenSCon (UnsafeFrom_ ctyp) = + \(e :* End) -> + \loc -> + do eE <- flattenWithName e + cE <- coerceFromCG ctyp eE + putExprStat $ loc .=. cE + +----------------------------------- +-- SCons in the Stochastic Monad -- +----------------------------------- + +flattenSCon (MeasureOp_ op) = flattenMeasureOp op + +flattenSCon Dirac = + \(e :* End) -> + \loc -> + do sE <- flattenWithName' e "samp" + putExprStat $ mdataWeight loc .=. (floatE 0) + putExprStat $ mdataSample loc .=. sE + +flattenSCon MBind = + \(ma :* b :* End) -> + \loc -> + caseBind b $ \v@(Variable _ _ typ) mb -> + do -- first + mE <- flattenWithName' ma "m" + + -- assign that sample to var + vId <- createIdent v + declare typ vId + assign vId (mdataSample mE) + flattenABT mb loc + putExprStat $ mdataWeight loc .+=. (mdataWeight mE) + +-- for now plats make use of a global sample +flattenSCon Plate = + \(size :* b :* End) -> + \loc -> + caseBind b $ \v@(Variable _ _ typ) body -> + do sizeE <- flattenWithName' size "s" + isMM <- managedMem <$> get + when (not isMM) (error "plate will leak memory without the '-g' flag and boehm-gc") + putExprStat $ (arraySize . mdataSample $ loc) .=. sizeE + putMallocStat (arrayData . mdataSample $ loc) sizeE (typeOf body) + + weightId <- genIdent' "w" + declare SProb weightId + let weightE = CVar weightId + assign weightId (floatE 0) + + itId <- createIdent v + declare SNat itId + let itE = CVar itId + currInd = index (arrayData . mdataSample $ loc) itE + + sampId <- genIdent' "samp" + declare (typeOf $ body) sampId + let sampE = CVar sampId + + reductionCG (Left CAddOp) + weightId + (itE .=. (intE 0)) + (itE .<. sizeE) + (CUnary CPostIncOp itE) + (do flattenABT body sampE + putExprStat (currInd .=. (mdataSample sampE)) + putExprStat (weightE .+=. (mdataWeight sampE))) + + putExprStat $ mdataWeight loc .=. weightE + +----------------------------------- +-- SCon's that arent implemented -- +----------------------------------- + +flattenSCon x = \_ -> \_ -> error $ "TODO: flattenSCon: " ++ show x + + + +-------------------------------------------------------------------------------- +-- NaryOps -- +-------------------------------------------------------------------------------- + +flattenNAryOp :: ABT Term abt + => NaryOp a + -> S.Seq (abt '[] a) + -> (CExpr -> CodeGen ()) +flattenNAryOp op args = + \loc -> + do es <- T.mapM flattenWithName args + case op of + And -> boolNaryOp op es loc + Or -> boolNaryOp op es loc + Xor -> boolNaryOp op es loc + Iff -> boolNaryOp op es loc + (Sum HSemiring_Prob) -> logSumExpCG es loc + _ -> let opE = F.foldr (binaryOp op) (S.index es 0) (S.drop 1 es) + in putExprStat (loc .=. opE) + + where boolNaryOp op' es' loc' = + let indexOf x = CMember x (Ident "index") True + es'' = fmap indexOf es' + expr = F.foldr (binaryOp op') + (S.index es'' 0) + (S.drop 1 es'') + in putExprStat ((indexOf loc') .=. expr) + + +-------------------------------------------------------------------------------- +-- Literals -- +-------------------------------------------------------------------------------- + +flattenLit + :: Literal a + -> (CExpr -> CodeGen ()) +flattenLit lit = + \loc -> + case lit of + (LNat x) -> putExprStat $ loc .=. (intE $ fromIntegral x) + (LInt x) -> putExprStat $ loc .=. (intE x) + (LReal x) -> putExprStat $ loc .=. (floatE $ fromRational x) + (LProb x) -> let rat = fromNonNegativeRational x + x' = (fromIntegral $ numerator rat) + / (fromIntegral $ denominator rat) + xE = log1pE (floatE x' .-. intE 1) + in putExprStat (loc .=. xE) + +-------------------------------------------------------------------------------- +-- Array and ArrayOps -- +-------------------------------------------------------------------------------- + + +flattenArray + :: (ABT Term abt) + => (abt '[] 'HNat) + -> (abt '[ 'HNat ] a) + -> (CExpr -> CodeGen ()) +flattenArray arity body = + \loc -> + caseBind body $ \v body' -> do + let arityE = arraySize loc + dataE = arrayData loc + typ = typeOf body' + + flattenABT arity arityE + + isManagedMem <- managedMem <$> get + let malloc' = if isManagedMem then gcMalloc else mallocE + putExprStat $ dataE + .=. (CCast (CTypeName (buildType typ) True) + (malloc' (arityE .*. (CSizeOfType (CTypeName (buildType typ) False))))) + + itId <- createIdent v + declare SNat itId + let itE = CVar itId + currInd = index dataE itE + + putStat $ opComment "Begin Array" + forCG (itE .=. (intE 0)) + (itE .<. arityE) + (CUnary CPostIncOp itE) + (flattenABT body' currInd) + putStat $ opComment "End Array" + +flattenArrayLiteral + :: ( ABT Term abt ) + => [abt '[] a] + -> (CExpr -> CodeGen ()) +flattenArrayLiteral es = + \loc -> do + arrId <- genIdent + isManagedMem <- managedMem <$> get + let arity = fromIntegral . length $ es + typ = typeOf . head $ es + arrE = CVar arrId + malloc' = if isManagedMem then gcMalloc else mallocE + + declare (SArray typ) arrId + putExprStat $ (arrayData arrE) + .=. (CCast (CTypeName (buildType typ) True) + (malloc' ((intE arity) .*. (CSizeOfType (CTypeName (buildType typ) False))))) + + putExprStat $ arraySize arrE .=. (intE arity) + sequence_ . snd $ foldl (\(i,acc) e -> (succ i,(assignIndex e i arrE):acc)) + (0,[]) + es + putExprStat $ loc .=. arrE + where assignIndex + :: ( ABT Term abt ) + => abt '[] a + -> Integer + -> (CExpr -> CodeGen ()) + assignIndex e index loc = do + eE <- flattenWithName e + putExprStat $ indirect ((arrayData loc) .+. (intE index)) .=. eE + +-------------- +-- ArrayOps -- +-------------- + +flattenArrayOp + :: ( ABT Term abt + , typs ~ UnLCs args + , args ~ LCs typs + ) + => ArrayOp typs a + -> SArgs abt args + -> (CExpr -> CodeGen ()) + + +flattenArrayOp (Index _) = + \(arr :* ind :* End) -> + \loc -> + do indE <- flattenWithName ind + arrE <- flattenWithName arr + let valE = index (CMember arrE (Ident "data") True) indE + putExprStat (loc .=. valE) + +flattenArrayOp (Size _) = + \(arr :* End) -> + \loc -> + do arrE <- flattenWithName arr + putExprStat (loc .=. (CMember arrE (Ident "size") True)) + +flattenArrayOp (Reduce _) = error "TODO: flattenArrayOp" + -- \(fun :* base :* arr :* End) -> + -- do funE <- flattenABT fun + -- baseE <- flattenABT base + -- arrE <- flattenABT arr + -- accI <- genIdent' "acc" + -- iterI <- genIdent' "iter" + + -- let sizeE = CMember arrE (Ident "size") True + -- iterE = CVar iterI + -- accE = CVar accI + -- cond = iterE .<. sizeE + -- inc = CUnary CPostIncOp iterE + + -- declare (typeOf base) accI + -- declare SInt iterI + -- assign accI baseE + -- forCG (iterE .=. (intE 0)) cond inc $ + -- assign accI $ CCall funE [accE] + + -- return accE + +-------------------------------------------------------------------------------- +-- Bucket and Reducers -- +-------------------------------------------------------------------------------- +{- Declarations for buckets - + since we will have some product of monoids we need unique names for each one. + + Ex: + bucket i from 0 to 100: + fanout(add(\_ -> 1),add(\_ -> 2)) + + we will need to keep track of two ints: + + int x; + int y; + for (i = 0; i < 100; i++) { + x += 1; + y += 2; + } + + Summary objects are nested pairs, e.g. pair(nat,pair(real,array(nat))) +-} + +flattenBucket + :: (ABT Term abt) + => abt '[] 'HNat + -> abt '[] 'HNat + -> Reducer abt '[] a + -> (CExpr -> CodeGen ()) +flattenBucket lo hi red = \loc -> do + putStat $ opComment "Begin Bucket" + loE <- flattenWithName' lo "lo" + hiE <- flattenWithName' hi "hi" + itId <- genIdent' "it" + declare SNat itId + let itE = CVar itId + initRed red loc + forCG (itE .=. loE) + (itE .<. hiE) + (CUnary CPostIncOp itE) + (accumRed red itE loc) + putStat $ opComment "End Bucket" + where initRed + :: (ABT Term abt) + => Reducer abt xs a + -> (CExpr -> CodeGen ()) + initRed mr = \loc -> + case mr of + (Red_Fanout mr1 mr2) -> initRed mr1 (datumFst loc) + >> initRed mr2 (datumSnd loc) + (Red_Split _ mr1 mr2) -> initRed mr1 (datumFst loc) + >> initRed mr2 (datumSnd loc) + (Red_Index s _ body) -> + let (vs,s') = caseBinds s + btyp = typeOfReducer body in + do sequence_ . foldMap11 + (\v' -> case v' of + (Variable _ _ typ') -> + [declare typ' =<< createIdent v']) + $ vs + sE <- flattenWithName s' + putExprStat $ arraySize loc .=. sE + putMallocStat (arrayData loc) sE btyp + itId <- genIdent + declare SNat itId + let itE = CVar itId + forCG (itE .=. (intE 0)) + (itE .<. sE) + (CUnary CPostIncOp itE) + (initRed body (index (arrayData loc) itE)) + Red_Nop -> return () + (Red_Add sr _) -> + putExprStat $ loc .=. (addMonoidIdentity . sing_HSemiring $ sr) + + accumRed + :: (ABT Term abt) + => Reducer abt xs a + -> CExpr + -> (CExpr -> CodeGen ()) + accumRed mr itE = \loc -> + case mr of + (Red_Index _ a body) -> + caseBind a $ \v@(Variable _ _ typ) a' -> + let (vs,a'') = caseBinds a' in + do vId <- createIdent v + declare typ vId + putExprStat $ (CVar vId) .=. itE + sequence_ . foldMap11 + (\v' -> case v' of + (Variable _ _ typ') -> + [declare typ' =<< createIdent v']) + $ vs + aE <- flattenWithName' a'' "index" + accumRed body itE (index (arrayData loc) aE) + (Red_Fanout mr1 mr2) -> accumRed mr1 itE (datumFst loc) + >> accumRed mr2 itE (datumSnd loc) + (Red_Split b mr1 mr2) -> + caseBind b $ \v@(Variable _ _ typ) b' -> + let (vs,b'') = caseBinds b' in + do vId <- createIdent v + declare typ vId + putExprStat $ (CVar vId) .=. itE + sequence_ . foldMap11 + (\v' -> case v' of + (Variable _ _ typ') -> + [declare typ' =<< createIdent v']) + $ vs + bE <- flattenWithName' b'' "cond" + ifCG (bE ... "index" .==. (intE 0)) + (accumRed mr1 itE (datumFst loc)) + (accumRed mr2 itE (datumSnd loc)) + Red_Nop -> return () + (Red_Add sr e) -> + caseBind e $ \v@(Variable _ _ typ) e' -> + let (vs,e'') = caseBinds e' in + do vId <- createIdent v + declare typ vId + putExprStat $ (CVar vId) .=. itE + sequence_ . foldMap11 + (\v' -> case v' of + (Variable _ _ typ') -> + [declare typ' =<< createIdent v']) + $ vs + eE <- flattenWithName e'' + case sing_HSemiring sr of + SProb -> logSumExpCG (S.fromList [loc,eE]) loc + _ -> putExprStat $ loc .+=. eE + +addMonoidIdentity :: Sing (a :: Hakaru) -> CExpr +addMonoidIdentity s = + case s of + SNat -> intE 0 + SInt -> intE 0 + SReal -> floatE 0 + SProb -> logE (floatE 0) + SArray x -> addMonoidIdentity x + x -> error $ "addMonoidIdentity{" ++ show x ++ "}" + +-------------------------------------------------------------------------------- +-- Datum and Case -- +-------------------------------------------------------------------------------- +{- + +Datum are sums of products of types. This maps to a C structure. flattenDatum +will produce a literal of some datum type. This will also produce a global +struct representing that datum which will be needed for the C compiler. + +-} + + +flattenDatum + :: (ABT Term abt) + => Datum (abt '[]) (HData' a) + -> (CExpr -> CodeGen ()) +flattenDatum (Datum _ typ code) = + \loc -> + do extDeclareTypes typ + assignDatum code loc + +assignDatum + :: (ABT Term abt) + => DatumCode xss (abt '[]) c + -> CExpr + -> CodeGen () +assignDatum code ident = + let index = getIndex code + indexExpr = CMember ident (Ident "index") True + in do putExprStat (indexExpr .=. (intE index)) + sequence_ $ assignSum code ident + where getIndex :: DatumCode xss b c -> Integer + getIndex (Inl _) = 0 + getIndex (Inr rest) = succ (getIndex rest) + +assignSum + :: (ABT Term abt) + => DatumCode xs (abt '[]) c + -> CExpr + -> [CodeGen ()] +assignSum code ident = fst $ runState (assignSum' code ident) cNameStream + +assignSum' + :: (ABT Term abt) + => DatumCode xs (abt '[]) c + -> CExpr + -> State [String] [CodeGen ()] +assignSum' (Inr rest) topIdent = + do (_:names) <- get + put names + assignSum' rest topIdent +assignSum' (Inl prod) topIdent = + do (name:_) <- get + return $ assignProd prod topIdent (CVar . Ident $ name) + +assignProd + :: (ABT Term abt) + => DatumStruct xs (abt '[]) c + -> CExpr + -> CExpr + -> [CodeGen ()] +assignProd dstruct topIdent sumIdent = + fst $ runState (assignProd' dstruct topIdent sumIdent) cNameStream + +assignProd' + :: (ABT Term abt) + => DatumStruct xs (abt '[]) c + -> CExpr + -> CExpr + -> State [String] [CodeGen ()] +assignProd' Done _ _ = return [] +assignProd' (Et (Konst d) rest) topIdent (CVar sumIdent) = + do (name:names) <- get + put names + let varName = CMember (CMember (CMember topIdent + (Ident "sum") + True) + sumIdent + True) + (Ident name) + True + rest' <- assignProd' rest topIdent (CVar sumIdent) + return $ [flattenABT d varName] ++ rest' +assignProd' _ _ _ = error $ "TODO: assignProd Ident" + + +---------- +-- Case -- +---------- + +-- currently we can only match on boolean values +flattenCase + :: forall abt a b + . (ABT Term abt) + => abt '[] a + -> [Branch a abt b] + -> (CExpr -> CodeGen ()) + +flattenCase c [ Branch (PDatum _ (PInl PDone)) trueB + , Branch (PDatum _ (PInr (PInl PDone))) falseB ] = + \loc -> + do cE <- flattenWithName c + ifCG ((cE ... "index") .==. (intE 0)) + (flattenABT trueB loc) + (flattenABT falseB loc) + +flattenCase e [ Branch (PDatum _ (PInl (PEt (PKonst PVar) + (PEt (PKonst PVar) + PDone)))) b + ] + = \loc -> do + eE <- flattenWithName e + caseBind b $ \vfst@(Variable _ _ fstTyp) b' -> + caseBind b' $ \vsnd@(Variable _ _ sndTyp) b'' -> do + fstId <- createIdent vfst + sndId <- createIdent vsnd + declare fstTyp fstId + declare sndTyp sndId + putExprStat $ (CVar fstId) .=. (datumFst eE) + putExprStat $ (CVar sndId) .=. (datumSnd eE) + flattenABT b'' loc + +flattenCase e _ = error $ "TODO: flattenCase{" ++ show (typeOf e) ++ "}" + + +-------------------------------------------------------------------------------- +-- PrimOp -- +-------------------------------------------------------------------------------- + +flattenPrimOp + :: ( ABT Term abt + , typs ~ UnLCs args + , args ~ LCs typs) + => PrimOp typs a + -> SArgs abt args + -> (CExpr -> CodeGen ()) + + +flattenPrimOp Pi = + \End -> + \loc -> let piE = log1pE ((CVar . Ident $ "M_PI") .-. (intE 1)) in + putExprStat (loc .=. piE) + +flattenPrimOp Not = + \(a :* End) -> + \loc -> + do aE <- flattenWithName a + bId <- genIdent + declare (typeOf a) bId + let datumIndex e = CMember e (Ident "index") True + bE = CVar bId + putExprStat $ datumIndex bE .=. (CCond (datumIndex aE .==. (intE 1)) + (intE 0) + (intE 1)) + putExprStat $ loc .=. bE + +flattenPrimOp RealPow = + \(base :* power :* End) -> + \loc -> + do baseE <- flattenWithName base + powerE <- flattenWithName power + let realPow = CCall (CVar . Ident $ "pow") + [ expm1E baseE .+. (intE 1), powerE] + putExprStat $ loc .=. (log1pE (realPow .-. (intE 1))) + +flattenPrimOp (NatPow baseTyp) = + \(base :* power :* End) -> + \loc -> + let sBase = sing_HSemiring baseTyp in + do baseId <- genIdent + declare sBase baseId + let baseE = CVar baseId + flattenABT base baseE + powerE <- flattenWithName power + let powerOf x y = CCall (CVar . Ident $ "pow") [x,y] + value = case sBase of + SProb -> log1pE $ (powerOf (expm1E baseE .+. (intE 1)) powerE) + .-. (intE 1) + _ -> powerOf baseE powerE + putExprStat $ loc .=. value + +flattenPrimOp (NatRoot baseTyp) = + \(base :* root :* End) -> + \loc -> + let sBase = sing_HRadical baseTyp in + do baseId <- genIdent + declare sBase baseId + let baseE = CVar baseId + flattenABT base baseE + rootE <- flattenWithName root + let powerOf x y = CCall (CVar . Ident $ "pow") [x,y] + recipE = (floatE 1) ./. rootE + value = case sBase of + SProb -> log1pE $ (powerOf (expm1E baseE .+. (intE 1)) recipE) + .-. (intE 1) + _ -> powerOf baseE recipE + putExprStat $ loc .=. value + +flattenPrimOp (Recip t) = + \(a :* End) -> + \loc -> + do aE <- flattenWithName a + case t of + HFractional_Real -> putExprStat $ loc .=. ((intE 1) ./. aE) + HFractional_Prob -> putExprStat $ loc .=. (CUnary CMinOp aE) + +-- | exp : real -> prob, because of this we can just turn it into a prob without taking +-- its log, which would give us an exp in the log-domain +flattenPrimOp Exp = \(a :* End) -> flattenABT a + +flattenPrimOp (Equal _) = + \(a :* b :* End) -> + \loc -> + do aE <- flattenWithName a + bE <- flattenWithName b + + -- special case for booleans + let aE' = case (typeOf a) of + (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CMember aE (Ident "index") True) + _ -> aE + let bE' = case (typeOf b) of + (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CMember bE (Ident "index") True) + _ -> bE + + putExprStat $ (CMember loc (Ident "index") True) + .=. (CCond (aE' .==. bE') (intE 0) (intE 1)) + + +flattenPrimOp (Less _) = + \(a :* b :* End) -> + \loc -> + do aE <- flattenWithName a + bE <- flattenWithName b + putExprStat $ (CMember loc (Ident "index") True) + .=. (CCond (aE .<. bE) (intE 0) (intE 1)) + +flattenPrimOp (Negate _) = + \(a :* End) -> + \loc -> + do aE <- flattenWithName a + putExprStat $ loc .=. (CUnary CMinOp $ aE) + +flattenPrimOp t = \_ -> error $ "TODO: flattenPrimOp: " ++ show t + + +-------------------------------------------------------------------------------- +-- MeasureOps and Superpose -- +-------------------------------------------------------------------------------- + +{- + +The sections contains operations in the stochastic monad. See also +(Dirac, MBind, and Plate) found in SCon. Also see Reject found at the top level. + +Remember in the C runtime. Measures are housed in a measure function, which +takes an `struct mdata` location. The MeasureOp attempts to store a value at +that location and returns 0 if it fails and 1 if it succeeds in that task. + +The functions uniformCG, normalCG, and gammaCG are primitives that will generate +functions and call them (similar to logSumExpCG). The reduce code size and make +samplers a little more readable. + +TODO: add inline pragmas to uniformCG, normalCG, and gammaCG +-} + +uniformFun :: CFunDef +uniformFun = CFunDef (CTypeSpec <$> retTyp) + (CDeclr Nothing (CDDeclrIdent funcId)) + [typeDeclaration SReal loId + ,typeDeclaration SReal hiId] + (CCompound . concat + $ [ CBlockDecl <$> [declMD] + , CBlockStat <$> comment ++ [assW,assS,CReturn . Just $ mE]] + ) + where r = castTo [CDouble] randE + rMax = castTo [CDouble] (CVar . Ident $ "RAND_MAX") + retTyp = buildType (SMeasure SReal) + (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) + (loId,loE) = let ident = Ident "lo" in (ident,CVar ident) + (hiId,hiE) = let ident = Ident "hi" in (ident,CVar ident) + value = (loE .+. ((r ./. rMax) .*. (hiE .-. loE))) + comment = fmap CComment + ["uniform :: real -> real -> *(mdata real) -> ()" + ,"------------------------------------------------"] + declMD = buildDeclaration (head retTyp) mId + assW = CExpr . Just $ mdataWeight mE .=. (floatE 0) + assS = CExpr . Just $ mdataSample mE .=. value + funcId = Ident "uniform" + + +uniformCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) +uniformCG aE bE = + \loc -> do + uId <- reserveIdent "uniform" + extDeclareTypes (SMeasure SReal) + extDeclare . CFunDefExt $ uniformFun + putExprStat $ loc .=. CCall (CVar uId) [aE,bE] + + +{- + This is very cryptic, but I assure you it is only building an AST for the + Marsaglia Polar Method +-} + +normalFun :: CFunDef +normalFun = CFunDef (CTypeSpec <$> retTyp) + (CDeclr Nothing (CDDeclrIdent (Ident "normal"))) + [typeDeclaration SReal aId + ,typeDeclaration SProb bId ] + ( CCompound . concat + $ [[CBlockDecl declMD],comment,decls,stmts]) + + where r = castTo [CDouble] randE + rMax = castTo [CDouble] (CVar . Ident $ "RAND_MAX") + retTyp = buildType (SMeasure SReal) + (aId,aE) = let ident = Ident "a" in (ident,CVar ident) + (bId,bE) = let ident = Ident "b" in (ident,CVar ident) + (qId,qE) = let ident = Ident "q" in (ident,CVar ident) + (uId,uE) = let ident = Ident "u" in (ident,CVar ident) + (vId,vE) = let ident = Ident "v" in (ident,CVar ident) + (rId,rE) = let ident = Ident "r" in (ident,CVar ident) + (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) + declMD = buildDeclaration (head retTyp) mId + draw xE = CExpr . Just $ xE .=. (((r ./. rMax) .*. (floatE 2)) .-. (floatE 1)) + body = seqCStat [draw uE + ,draw vE + ,CExpr . Just $ qE .=. ((uE .*. uE) .+. (vE .*. vE))] + polar = CWhile (qE .>. (floatE 1)) body True + setR = CExpr . Just $ rE .=. (sqrtE (((CUnary CMinOp (floatE 2)) .*. logE qE) ./. qE)) + finalValue = aE .+. (uE .*. rE .*. bE) + comment = fmap (CBlockStat . CComment) + ["normal :: real -> real -> *(mdata real) -> ()" + ,"Marsaglia Polar Method" + ,"-----------------------------------------------"] + decls = (CBlockDecl . typeDeclaration SReal) <$> [uId,vId,qId,rId] + stmts = CBlockStat <$> [polar,setR, assW, assS,CReturn . Just $ mE] + assW = CExpr . Just $ mdataWeight mE .=. (floatE 0) + assS = CExpr . Just $ mdataSample mE .=. finalValue + + +normalCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) +normalCG aE bE = + \loc -> do + nId <- reserveIdent "normal" + extDeclareTypes (SMeasure SReal) + extDeclare . CFunDefExt $ normalFun + putExprStat $ loc .=. (CCall (CVar nId) [aE,bE]) + +{- + This method is from Marsaglia and Tsang "a simple method for generating gamma variables" +-} +gammaFun :: CFunDef +gammaFun = CFunDef (CTypeSpec <$> retTyp) + (CDeclr Nothing (CDDeclrIdent (Ident "gamma"))) + [typeDeclaration SProb aId + ,typeDeclaration SProb bId] + ( CCompound . concat + $ [[CBlockDecl declMD],comment,decls,stmts]) + where (aId,aE) = let ident = Ident "a" in (ident,CVar ident) + (bId,bE) = let ident = Ident "b" in (ident,CVar ident) + (cId,cE) = let ident = Ident "c" in (ident,CVar ident) + (dId,dE) = let ident = Ident "d" in (ident,CVar ident) + (xId,xE) = let ident = Ident "x" in (ident,CVar ident) + (vId,vE) = let ident = Ident "v" in (ident,CVar ident) + (uId,uE) = let ident = Ident "u" in (ident,CVar ident) + (mId,mE) = let ident = Ident "mdata" in (ident,CVar ident) + retTyp = buildType (SMeasure SProb) + declMD = buildDeclaration (head retTyp) mId + comment = fmap (CBlockStat . CComment) + ["gamma :: real -> prob -> *(mdata prob) -> ()" + ,"Marsaglia and Tsang 'a simple method for generating gamma variables'" + ,"--------------------------------------------------------------------"] + decls = fmap CBlockDecl $ (fmap (typeDeclaration SReal) [dId,cId,vId]) + ++ (fmap (typeDeclaration (SMeasure SReal)) [uId,xId]) + stmts = fmap CBlockStat $ [assD,assC,outerWhile] + xS = mdataSample xE + uS = mdataSample uE + assD = CExpr . Just $ dE .=. (aE .-. ((floatE 1) ./. (floatE 3))) + assC = CExpr . Just $ cE .=. ((floatE 1) ./. (sqrtE ((floatE 9) .*. dE))) + outerWhile = CWhile (intE 1) (seqCStat [innerWhile,assV,assU,exit]) False + innerWhile = CWhile (vE .<=. (floatE 0)) (seqCStat [assX,assVIn]) True + assX = CExpr . Just $ xE .=. (CCall (CVar . Ident $ "normal") [(floatE 0),(floatE 1)]) + assVIn = CExpr . Just $ vE .=. ((floatE 1) .+. (cE .*. xS)) + assV = CExpr . Just $ vE .=. (vE .*. vE .*. vE) + assU = CExpr . Just $ uE .=. (CCall (CVar . Ident $ "uniform") [(floatE 0),(floatE 1)]) + exitC1 = uS .<. ((floatE 1) .-. ((floatE 0.331 .*. (xS .*. xS) .*. (xS .*. xS)))) + exitC2 = (logE uS) .<. (((floatE 0.5) .*. (xS .*. xS)) .+. (dE .*. ((floatE 1.0) .-. vE .+. (logE vE)))) + assW = CExpr . Just $ mdataWeight mE .=. (floatE 0) + assS = CExpr . Just $ mdataSample mE .=. (logE (dE .*. vE)) .+. bE + exit = CIf (exitC1 .||. exitC2) (seqCStat [assW,assS,CReturn . Just $ mE]) Nothing + + +gammaCG :: CExpr -> CExpr -> (CExpr -> CodeGen ()) +gammaCG aE bE = + \loc -> do + extDeclareTypes (SMeasure SReal) + (_:_:gId:[]) <- mapM reserveIdent ["uniform","normal","gamma"] + mapM_ (extDeclare . CFunDefExt) [uniformFun,normalFun,gammaFun] + putExprStat $ loc .=. (CCall (CVar gId) [aE,bE]) + + +flattenMeasureOp + :: forall abt typs args a . + ( ABT Term abt + , typs ~ UnLCs args + , args ~ LCs typs ) + => MeasureOp typs a + -> SArgs abt args + -> (CExpr -> CodeGen ()) + + +flattenMeasureOp Uniform = + \(a :* b :* End) -> + \loc -> + do aE <- flattenWithName a + bE <- flattenWithName b + uniformCG aE bE loc + + +flattenMeasureOp Normal = + \(a :* b :* End) -> + \loc -> + do aE <- flattenWithName a + bE <- flattenWithName b + normalCG aE (expE bE) loc + +flattenMeasureOp Poisson = + \(lam :* End) -> + \loc -> + do lamE <- flattenWithName lam + (lId:kId:pId:[]) <- mapM genIdent' ["l","k","p"] + declare SProb lId + declare SNat kId + declare SProb pId + let (lE:kE:pE:[]) = fmap CVar [lId,kId,pId] + putExprStat $ lE .=. (expE (CUnary CMinOp $ expE lamE)) + putExprStat $ kE .=. (intE 0) + putExprStat $ pE .=. (floatE 1) + doWhileCG (pE .>. lE) $ + do uId <- genIdent' "u" + declare (SMeasure SReal) uId + let uE = CVar uId + uniformCG (intE 0) (intE 1) uE + putExprStat $ pE .*=. (mdataSample uE) + putExprStat $ kE .+=. (intE 1) + + putExprStat $ mdataWeight loc .=. (intE 0) + putExprStat $ mdataSample loc .=. (kE .-. (intE 1)) + +flattenMeasureOp Gamma = + \(a :* b :* End) -> + \loc -> + do aE <- flattenWithName a + bE <- flattenWithName b + gammaCG (expE aE) bE loc + + +flattenMeasureOp Beta = + \(a :* b :* End) -> flattenABT (HKP.beta'' a b) + + +-- I ran into a bug here where sometime I recieved a location by reference and +-- others by value. Since measureOps assign a sample to mdata that they have a +-- reference to, we should enforce that when passing around mdata it is by +-- reference +flattenMeasureOp Categorical = \(arr :* End) -> + \loc -> + do arrE <- flattenWithName arr + + itId <- genIdent' "it" + declare SNat itId + let itE = CVar itId + + -- Accumulator for the total probability of the input array + wSumId <- genIdent' "ws" + declare SProb wSumId + let wSumE = CVar wSumId + assign wSumId (logE (intE 0)) + + -- Accumulator for the max value in the input array + wMaxId <- genIdent' "max" + declare SProb wMaxId + let wMaxE = CVar wMaxId + assign wMaxId (logE (floatE 0)) + + let currE = index (arrayData arrE) itE + + isPar <- isParallel + mkSequential + + -- Calculate the maximum value of the input array + -- And calculate the total weight + forCG (itE .=. (intE 0)) + (itE .<. (arraySize arrE)) + (CUnary CPostIncOp itE) $ do + ifCG (wMaxE .<. currE) + (putExprStat $ wMaxE .=. currE) + (return ()) + logSumExpCG (S.fromList [wSumE, currE]) wSumE + putExprStat $ wSumE .=. (wSumE .-. wMaxE) + + -- draw number from uniform(0, weightSum) + rId <- genIdent' "r" + declare SReal rId + let r = castTo [CDouble] randE + rMax = castTo [CDouble] (CVar . Ident $ "RAND_MAX") + rE = CVar rId + assign rId (logE (r ./. rMax) .+. wSumE) + + assign wSumId (logE (intE 0)) + assign itId (intE 0) + whileCG (intE 1) $ + do ifCG (rE .<. wSumE) + (do putExprStat $ mdataWeight loc .=. (intE 0) + putExprStat $ mdataSample loc .=. (itE .-. (intE 1)) + putStat CBreak) + (return ()) + logSumExpCG (S.fromList [wSumE, currE .-. wMaxE]) wSumE + putExprStat $ CUnary CPostIncOp itE + + when isPar mkParallel + + +flattenMeasureOp x = error $ "TODO: flattenMeasureOp: " ++ show x + +--------------- +-- Superpose -- +--------------- + +flattenSuperpose + :: (ABT Term abt) + => NE.NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a)) + -> (CExpr -> CodeGen ()) + +-- do we need to normalize? +flattenSuperpose pairs = + let pairs' = NE.toList pairs in + + if length pairs' == 1 + then \loc -> let (w,m) = head pairs' in + do mE <- flattenWithName m + wE <- flattenWithName w + putExprStat $ mdataWeight loc .=. ((mdataWeight mE) .+. wE) + putExprStat $ mdataSample loc .=. (mdataSample mE) + + else \loc -> + do wEs <- mapM (\(w,_) -> flattenWithName' w "w") pairs' + + wSumId <- genIdent' "ws" + declare SProb wSumId + let wSumE = CVar wSumId + logSumExpCG (S.fromList wEs) wSumE + + -- draw number from uniform(0, weightSum) + rId <- genIdent' "r" + declare SReal rId + let r = castTo [CDouble] randE + rMax = castTo [CDouble] (CVar . Ident $ "RAND_MAX") + rE = CVar rId + assign rId ((r ./. rMax) .*. (expE wSumE)) + + -- an iterator for picking a measure + itId <- genIdent' "it" + declare SProb itId + let itE = CVar itId + assign itId (logE (intE 0)) + + -- an output measure to assign to + outId <- genIdent' "out" + declare (typeOf . snd . head $ pairs') outId + let outE = CVar outId + + outLabel <- genIdent' "exit" + + forM_ (zip wEs pairs') + $ \(wE,(_,m)) -> + do logSumExpCG (S.fromList [itE,wE]) itE + ifCG (rE .<. (expE itE)) + (flattenABT m outE >> putStat (CGoto outLabel)) + (return ()) + + putStat $ CLabel outLabel (CExpr Nothing) + putExprStat $ mdataWeight loc .=. ((mdataWeight outE) .+. wSumE) + putExprStat $ mdataSample loc .=. (mdataSample outE) + +-------------------------------------------------------------------------------- +-- Specialized Arithmetic -- +-------------------------------------------------------------------------------- + +-------------------------------------- +-- LogSumExp for NaryOp Add [SProb] -- +-------------------------------------- +{- + +Special for addition of probabilities we have a logSumExp. This will compute the +sum of the probabilities safely. Just adding the exp(a . prob) would make us +loose any of the safety from underflow that we got from storing prob in the log +domain + +-} + +-- the tree traversal is a depth first search +logSumExp :: S.Seq CExpr -> CExpr +logSumExp es = mkCompTree 0 1 + + where lastIndex = S.length es - 1 + + compIndices :: Int -> Int -> CExpr -> CExpr -> CExpr + compIndices i j = CCond ((S.index es i) .>. (S.index es j)) + + mkCompTree :: Int -> Int -> CExpr + mkCompTree i j + | j == lastIndex = compIndices i j (logSumExp' i) (logSumExp' j) + | otherwise = compIndices i j + (mkCompTree i (succ j)) + (mkCompTree j (succ j)) + + diffExp :: Int -> Int -> CExpr + diffExp a b = expm1E ((S.index es a) .-. (S.index es b)) + + -- given the max index, produce a logSumExp expression + logSumExp' :: Int -> CExpr + logSumExp' 0 = S.index es 0 + .+. (log1pE $ foldr (\x acc -> diffExp x 0 .+. acc) + (diffExp 1 0) + [2..S.length es - 1] + .+. (intE $ fromIntegral lastIndex)) + logSumExp' i = S.index es i + .+. (log1pE $ foldr (\x acc -> if i == x + then acc + else diffExp x i .+. acc) + (diffExp 0 i) + [1..S.length es - 1] + .+. (intE $ fromIntegral lastIndex)) + + +-- | logSumExpCG creates global functions for every n-ary logSumExp function +-- this reduces code size +logSumExpCG :: S.Seq CExpr -> (CExpr -> CodeGen ()) +logSumExpCG seqE = + let size = S.length $ seqE + name = "logSumExp" ++ (show size) + funcId = Ident name + in \loc -> do -- reset the names so that the function is the same for each arity + let argIds = fmap Ident (take size cNameStream) + decls = fmap (typeDeclaration SProb) argIds + vars = fmap CVar argIds + funCG CDouble funcId decls + (putStat . CReturn . Just . logSumExp . S.fromList $ vars) + putExprStat $ loc .=. (CCall (CVar funcId) (F.toList seqE)) + +------------------------------------- +-- LogSumExp for Summation of Prob -- +------------------------------------- +{- + +For summation of SProb we need a new logSumExp function that will find the max +of an array and then sum it in a loop + +-} + +lseSummateArrayCG + :: ( ABT Term abt ) + => (abt '[ a ] b) + -> CExpr + -> (CExpr -> CodeGen ()) +lseSummateArrayCG body arrayE = + caseBind body $ \v body' -> + \loc -> do + (maxVId:maxIId:sumId:[]) <- mapM genIdent' ["maxV","maxI","sum"] + itId <- createIdent v + mapM_ (declare SProb) [maxVId,sumId] + mapM_ (declare SNat) [maxIId,itId] + let (maxVE:maxIE:sumE:itE:[]) = fmap CVar [maxVId,maxIId,sumId,itId] + forCG (itE .=. intE 0) + (itE .<. arraySize arrayE) + (CUnary CPostIncOp itE) + (do tmpId <- genIdent + declare SProb tmpId + let tmpE = CVar tmpId + flattenABT body' tmpE + putExprStat $ derefIndex itE .=. tmpE + putStat $ CIf ((maxVE .<. tmpE) .||. (itE .==. (intE 0))) + (seqCStat . fmap (CExpr . Just) $ + [ maxVE .=. tmpE + , maxIE .=. itE ]) + Nothing) + putExprStat $ sumE .=. (floatE 0) -- the sum is actually in real space + forCG (itE .=. intE 0) + (itE .<. arraySize arrayE) + (CUnary CPostIncOp itE) + (putStat $ CIf (itE .!=. maxIE) + (CExpr . Just $ sumE .+=. (expE ((derefIndex itE) .-. (maxVE)))) + Nothing) + + putExprStat $ loc .=. (maxVE .+. (log1pE sumE)) + + where derefIndex xE = index (arrayData arrayE) xE + +--------------------- +-- Kahan Summation -- +--------------------- +-- | given a body and a size compute the kahan summation. This should work on +-- both probs and reals +kahanSummationCG + :: ( ABT Term abt ) + => (abt '[ a ] b) + -> CExpr + -> CExpr + -> (CExpr -> CodeGen ()) +kahanSummationCG body loE hiE = + caseBind body $ \v body' -> + \loc -> do + (tId:cId:[]) <- mapM genIdent' ["t","c"] + itId <- createIdent v + declare SNat itId + mapM_ (declare SProb) [tId,cId] + let (tE:cE:itE:[]) = fmap CVar [tId,cId,itId] + putExprStat $ tE .=. (floatE 0) + putExprStat $ cE .=. (floatE 0) + forCG (itE .=. loE) + (itE .<. hiE) + (CUnary CPostIncOp itE) + (do (xId:yId:zId:[]) <- mapM genIdent' ["x","y","z"] + mapM_ (declare SProb) [xId,yId,zId] + let (xE:yE:zE:[]) = fmap CVar [xId,yId,zId] + flattenABT body' xE + putExprStat $ yE .=. (xE .-. cE) + putExprStat $ zE .=. (tE .+. yE) + putExprStat $ cE .=. ((zE .-. tE) .-. yE) + putExprStat $ tE .=. zE) + putExprStat $ loc .=. tE + +-------------------------------------------------------------------------------- +-- Coercion Helpers -- +-------------------------------------------------------------------------------- + +-- instance PrimCoerce Value where +-- primCoerceTo c l = +-- case (c,l) of +-- (Signed HRing_Int, VNat a) -> VInt $ fromNat a +-- (Signed HRing_Real, VProb a) -> VReal $ LF.fromLogFloat a +-- (Continuous HContinuous_Prob, VNat a) -> +-- VProb $ LF.logFloat (fromIntegral (fromNat a) :: Double) +-- (Continuous HContinuous_Real, VInt a) -> VReal $ fromIntegral a +-- _ -> error "no a defined primitive coercion" + +-- primCoerceFrom c l = +-- case (c,l) of +-- (Signed HRing_Int, VInt a) -> VNat $ unsafeNat a +-- (Signed HRing_Real, VReal a) -> VProb $ LF.logFloat a +-- (Continuous HContinuous_Prob, VProb a) -> +-- VNat $ unsafeNat $ floor (LF.fromLogFloat a :: Double) +-- (Continuous HContinuous_Real, VReal a) -> VInt $ floor a +-- _ -> error "no a defined primitive coercion" + + + +coerceToCG + :: forall (a :: Hakaru) (b :: Hakaru) + . Coercion a b + -> CExpr + -> CodeGen CExpr +coerceToCG (CCons (Signed HRing_Int) cs) e = nat2int e >>= coerceToCG cs +coerceToCG (CCons (Signed HRing_Real) cs) e = prob2real e >>= coerceToCG cs +coerceToCG (CCons (Continuous HContinuous_Prob) cs) e = nat2prob e >>= coerceToCG cs +coerceToCG (CCons (Continuous HContinuous_Real) cs) e = int2real e >>= coerceToCG cs +coerceToCG CNil e = return e + +coerceFromCG + :: forall (a :: Hakaru) (b :: Hakaru) + . Coercion a b + -> CExpr + -> CodeGen CExpr +coerceFromCG (CCons (Signed HRing_Int) cs) e = int2nat e >>= coerceFromCG cs +coerceFromCG (CCons (Signed HRing_Real) cs) e = real2prob e >>= coerceFromCG cs +coerceFromCG (CCons (Continuous HContinuous_Prob) cs) e = prob2nat e >>= coerceFromCG cs +coerceFromCG (CCons (Continuous HContinuous_Real) cs) e = real2int e >>= coerceFromCG cs +coerceFromCG CNil e = return e + +-- safe +nat2int,nat2prob,prob2real,int2real + :: CExpr -> CodeGen CExpr +nat2int x = return x +nat2prob x = do x' <- localVar' SProb "n2p" + putExprStat $ x' .=. (logE x) + return x' +prob2real x = do x' <- localVar' SProb "p2r" + putExprStat $ x' .=. ((expm1E x) .+. (floatE 1)) + return x' +int2real x = return (castTo [CDouble] x) + +-- unsafe +{- Because of the hkc representation of reals and probs as doubles, (instead of + rationals). we will just silently truncate values for prob2nat and real2int +-} +int2nat,prob2nat,real2prob,real2int + :: CExpr -> CodeGen CExpr +int2nat x = + do x' <- localVar' SNat "i2n" + ifCG (x .<. (intE 0)) + (do putExprStat $ printfE [ stringE "error: cannot coerce negative int to nat\n" ] + putExprStat $ mkCallE "abort" [] ) + (putExprStat $ x' .=. (castTo [CUnsigned, CInt] x)) + return x' +prob2nat x = return (castTo [CUnsigned, CInt] x) +real2prob x = + do x' <- localVar' SProb "r2p" + ifCG (x .<. (intE 0)) + (do putExprStat $ printfE [ stringE "error: cannot coerce negative real to prob\n" ] + putExprStat $ mkCallE "abort" [] ) + (putExprStat $ x' .=. (logE x)) + return x' +real2int x = return (castTo [CInt] x) diff --git a/haskell/Language/Hakaru/CodeGen/Libs.hs b/haskell/Language/Hakaru/CodeGen/Libs.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/CodeGen/Libs.hs @@ -0,0 +1,141 @@ +---------------------------------------------------------------- +-- 2016.12.20 +-- | +-- Module : Language.Hakaru.CodeGen.Libraries +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : zsulliva@indiana.edu +-- Stability : experimental +-- Portability : GHC-only +-- +-- Bindings to common C libraries +-- +---------------------------------------------------------------- + +module Language.Hakaru.CodeGen.Libs + ( -- math.h + expE, expm1E, logE, log1pE, sqrtE, + infinityE,negInfinityE, + + -- stdio.h + printfE, sscanfE, fopenE, fcloseE, fileT, feofE, fgetsE, rewindE, + + -- stdlib.h + randE, srandE, mallocE, freeE, + + -- Boehm Gargbage Collector + gcHeader, gcInit, gcMalloc, + + -- OpenMP + openMpHeader, ompGetNumThreads, ompGetThreadNum, + ) where + +import Language.Hakaru.CodeGen.AST + +{- + + As a convention to make the CExpressions standout, functions that return CExpr + have a suffix 'E' for instance 'printfE' + +-} + + +-------------------------------------------------------------------------------- +-- Lib C -- +-------------------------------------------------------------------------------- +{- + Here we have calls to a very small subset of functionality provided by libc. + In the future, we should have a standard way to add in bindings to C + libraries. Easily generating code for existing C libraries is one of the key + design goals of pedantic-c +-} + +------------ +-- math.h -- +------------ + +expE,expm1E,logE,log1pE,sqrtE :: CExpr -> CExpr +expE = mkUnaryE "exp" +expm1E = mkUnaryE "expm1" +logE = mkUnaryE "log" +log1pE = mkUnaryE "log1p" +sqrtE = mkUnaryE "sqrt" + +infinityE,negInfinityE :: CExpr +infinityE = (intE 1) ./. (intE 0) +negInfinityE = logE (intE 0) + +-------------- +-- stdlib.h -- +-------------- + +randE :: CExpr +randE = mkCallE "rand" [] + +srandE :: CExpr -> CExpr +srandE e = mkCallE "srand" [e] + +mallocE :: CExpr -> CExpr +mallocE = mkUnaryE "malloc" + +freeE :: CExpr -> CExpr +freeE = mkUnaryE "free" + +-------------- +-- stdio.h -- +-------------- + +printfE,sscanfE :: [CExpr] -> CExpr +printfE = mkCallE "printf" +sscanfE = mkCallE "sscanf" + +fopenE :: CExpr -> CExpr -> CExpr +fopenE e0 e1 = mkCallE "fopen" [e0,e1] + +fcloseE,feofE,rewindE :: CExpr -> CExpr +fcloseE e = mkCallE "fclose" [e] +feofE e = mkCallE "feof" [e] +rewindE e = mkCallE "rewind" [e] + +fgetsE :: CExpr -> CExpr -> CExpr -> CExpr +fgetsE e0 e1 e2 = mkCallE "fgets" [e0,e1,e2] + +fileT :: CTypeSpec +fileT = CTypeDefType (Ident "FILE") + +-------------------------------------------------------------------------------- +-- Boehm Garbage Collector -- +-------------------------------------------------------------------------------- +{- + Currently needed for handling arrays and datum. + + In the future, an intermediate language based on the region calculus will be + employed here. +-} + +gcHeader :: Preprocessor +gcHeader = PPInclude "gc.h" + +gcInit :: CExpr +gcInit = mkCallE "GC_INIT" [] + +gcMalloc :: CExpr -> CExpr +gcMalloc e = mkCallE "GC_MALLOC" [e] + +-------------------------------------------------------------------------------- +-- OpenMP -- +-------------------------------------------------------------------------------- +{- + For generating pragmas for shared memory parallelism, that is parallelism on + on a single process that makes use of multithreaded processors. This + interface is implemented in most C compilers and is accessed through pragmas +-} + +openMpHeader :: Preprocessor +openMpHeader = PPInclude "omp.h" + +ompGetNumThreads :: CExpr +ompGetNumThreads = mkCallE "omp_get_num_threads" [] + +ompGetThreadNum :: CExpr +ompGetThreadNum = mkCallE "omp_get_thread_num" [] diff --git a/haskell/Language/Hakaru/CodeGen/Pretty.hs b/haskell/Language/Hakaru/CodeGen/Pretty.hs --- a/haskell/Language/Hakaru/CodeGen/Pretty.hs +++ b/haskell/Language/Hakaru/CodeGen/Pretty.hs @@ -46,6 +46,11 @@ newline :: Doc newline = char '\n' +instance Pretty a => Pretty (Maybe a) where + pretty Nothing = empty + pretty (Just x) = pretty x + + -------------------------------------------------------------------------------- -- Top Level -- -------------------------------------------------------------------------------- @@ -99,8 +104,8 @@ declarators (dr, Just ilist) = pretty dr <+> text "=" <+> pretty ilist instance Pretty CDeclr where - pretty (CDeclr mp dds) = - mpretty mp <+> (hsep . fmap pretty $ dds) + pretty (CDeclr mp dd) = + mpretty mp <+> (pretty $ dd) instance Pretty CPtrDeclr where pretty (CPtrDeclr ts) = text "*" <+> (hsep . fmap pretty $ ts) @@ -109,7 +114,8 @@ pretty (CDDeclrIdent i) = pretty i pretty (CDDeclrArr dd e) = pretty dd <+> (brackets . pretty $ e) pretty (CDDeclrFun dd ts) = - pretty dd <+> (hsep . punctuate comma . fmap pretty $ ts) + pretty dd <> (parens . hsep . punctuate comma . fmap pretty $ ts) + pretty (CDDeclrRec declr) = parens . pretty $ declr instance Pretty CDeclSpec where @@ -140,9 +146,16 @@ pretty CSigned = text "signed" pretty CUnsigned = text "unsigned" pretty (CSUType cs) = pretty cs - pretty (CTypeDefType _) = error "TODO: Pretty TypeDef" + pretty (CTypeDefType sid) = pretty sid pretty (CEnumType _) = error "TODO: Pretty EnumType" +instance Pretty CTypeName where + pretty (CTypeName tspecs pb) = + let ss = sep . fmap pretty $ tspecs + in if pb + then ss <+> text "*" + else ss + instance Pretty CSUSpec where pretty (CSUSpec tag mi []) = pretty tag <+> mpretty mi @@ -290,6 +303,6 @@ instance Pretty CConst where pretty (CIntConst i) = text . show $ i - pretty (CCharConst c) = char c + pretty (CCharConst c) = text . show $ c pretty (CFloatConst f) = float f pretty (CStringConst s) = text . show $ s diff --git a/haskell/Language/Hakaru/CodeGen/Types.hs b/haskell/Language/Hakaru/CodeGen/Types.hs --- a/haskell/Language/Hakaru/CodeGen/Types.hs +++ b/haskell/Language/Hakaru/CodeGen/Types.hs @@ -19,15 +19,16 @@ module Language.Hakaru.CodeGen.Types ( buildDeclaration + , buildDeclaration' , buildPtrDeclaration -- tools for building C types , typeDeclaration , typePtrDeclaration + , typeName -- arrays , arrayDeclaration - , arrayName , arrayStruct , arraySize , arrayData @@ -37,7 +38,6 @@ -- mdata , mdataDeclaration , mdataPtrDeclaration - , mdataName , mdataStruct , mdataStruct' , mdataWeight @@ -45,26 +45,24 @@ , mdataPtrWeight , mdataPtrSample + -- datum , datumDeclaration - , datumName , datumStruct - , functionDef , datumSum , datumProd + , datumFst + , datumSnd + -- functions and closures + , functionDef + , closureDeclaration + , buildType - , mkDecl - , mkPtrDecl + , castTo + , castToPtrOf + , callStruct , buildStruct , buildUnion - - , intDecl - , natDecl - , doubleDecl - , doublePtr - , intPtr - , natPtr - , boolTyp , binaryOp ) where @@ -75,59 +73,100 @@ import Language.Hakaru.Types.HClasses import Language.Hakaru.Types.Sing import Language.Hakaru.CodeGen.AST - -import Prelude hiding (exp,log,sqrt) +import Language.Hakaru.CodeGen.Libs buildDeclaration :: CTypeSpec -> Ident -> CDecl buildDeclaration ctyp ident = CDecl [ CTypeSpec ctyp ] - [( CDeclr Nothing [ CDDeclrIdent ident ] + [( CDeclr Nothing (CDDeclrIdent ident) , Nothing)] +buildDeclaration' :: [CTypeSpec] -> Ident -> CDecl +buildDeclaration' specs ident = + CDecl (fmap CTypeSpec specs) + [( CDeclr Nothing (CDDeclrIdent ident) + , Nothing)] + buildPtrDeclaration :: CTypeSpec -> Ident -> CDecl buildPtrDeclaration ctyp ident = CDecl [ CTypeSpec ctyp ] - [( CDeclr (Just $ CPtrDeclr []) [ CDDeclrIdent ident ] + [( CDeclr (Just $ CPtrDeclr []) (CDDeclrIdent ident) , Nothing)] typeDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl typeDeclaration typ ident = CDecl (fmap CTypeSpec $ buildType typ) - [( CDeclr Nothing [ CDDeclrIdent ident ] + [( CDeclr Nothing (CDDeclrIdent ident) , Nothing)] typePtrDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl typePtrDeclaration typ ident = CDecl (fmap CTypeSpec $ buildType typ) [( CDeclr (Just $ CPtrDeclr []) - [ CDDeclrIdent ident ] + (CDDeclrIdent ident) , Nothing)] + +---------------- +-- Type Names -- +---------------- +{- +Type names are used when constructing C structures. In most cases there is a +unique C structure name for a Hakaru type. This is not the case for functions +which are compiled into closures, which are unique to a certain context and +Hakaru type. +-} + +typeName :: Sing (a :: Hakaru) -> String +typeName SInt = "int" +typeName SNat = "nat" +typeName SReal = "real" +typeName SProb = "prob" +typeName (SArray t) = "array_" ++ typeName t +typeName (SMeasure t) = "mdata_" ++ typeName t +typeName f@(SFun _ _) = error $ "typeName{SFun} doen't make sense: unknown context for {" ++ show f ++ "}" +typeName (SData _ t) = "dat_" ++ datumSumName t + where datumSumName :: Sing (a :: [[HakaruFun]]) -> String + datumSumName SVoid = "V" + datumSumName (SPlus p s) = datumProdName p ++ datumSumName s + + datumProdName :: Sing (a :: [HakaruFun]) -> String + datumProdName SDone = "D" + datumProdName (SEt x p) = datumPrimName x ++ datumProdName p + + datumPrimName :: Sing (a :: HakaruFun) -> String + datumPrimName SIdent = "I" + datumPrimName (SKonst s) = "K" ++ typeName s + + + + -------------------------------------------------------------------------------- --- Representing Hakaru Arrays +-- Arrays -- +-------------------------------------------------------------------------------- +{- + We represent arrays as structs with an 'unsigned int' for the size and a + pointer to a block of array elements. -arrayName :: Sing (a :: Hakaru) -> String -arrayName SInt = "arrayInt" -arrayName SNat = "arrayNat" -arrayName SReal = "arrayReal" -arrayName SProb = "arrayProb" -arrayName t = error $ "arrayName: cannot make array from type: " ++ show t + Because arrays may point to undeclared types (such as arrays of datum), we + need to return a list of external declarations with our array type +-} arrayStruct :: Sing (a :: Hakaru) -> CExtDecl arrayStruct t = CDeclExt (CDecl [CTypeSpec $ arrayStruct' t] []) arrayStruct' :: Sing (a :: Hakaru) -> CTypeSpec arrayStruct' t = aStruct - where aSize = buildDeclaration CInt (Ident "size") + where aSize = buildDeclaration' [CUnsigned,CInt] (Ident "size") aData = typePtrDeclaration t (Ident "data") - aStruct = buildStruct (Just . Ident . arrayName $ t) [aSize,aData] + aStruct = buildStruct (Just . Ident . typeName . SArray $ t) [aSize,aData] arrayDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl -arrayDeclaration typ = buildDeclaration (callStruct (arrayName typ)) +arrayDeclaration = buildDeclaration . callStruct . typeName . SArray arraySize :: CExpr -> CExpr @@ -142,20 +181,17 @@ arrayPtrData :: CExpr -> CExpr arrayPtrData e = CMember e (Ident "data") False --------------------------------------------------------------------------------- --- Measure Data -mdataName :: Sing (a :: Hakaru) -> String -mdataName SInt = "mdataInt" -mdataName SNat = "mdataNat" -mdataName SReal = "mdataReal" -mdataName SProb = "mdataProb" -mdataName (SArray SInt) = "mdataArrayInt" -mdataName (SArray SNat) = "mdataArrayNat" -mdataName (SArray SReal) = "mdataArrayReal" -mdataName (SArray SProb) = "mdataArrayProb" -mdataName t = error $ "mdataName: cannot make mdata from type: " ++ show t +-------------------------------------------------------------------------------- +-- Measure Data -- +-------------------------------------------------------------------------------- +{- + Measure datum are structures that will be used for sampling. We represent it + as a structure with a 'double' in log-domain corresponding to the weight of + the sample and an item of the sample type. +-} + mdataStruct :: Sing (a :: Hakaru) -> CExtDecl mdataStruct t = CDeclExt (CDecl [CTypeSpec $ mdataStruct' t] []) @@ -163,19 +199,19 @@ mdataStruct' t = mdStruct where weight = buildDeclaration CDouble (Ident "weight") sample = typeDeclaration t (Ident "sample") - mdStruct = buildStruct (Just . Ident . mdataName $ t) [weight,sample] + mdStruct = buildStruct (Just . Ident . typeName . SMeasure $ t) [weight,sample] mdataDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl -mdataDeclaration typ = buildDeclaration (callStruct (mdataName typ)) +mdataDeclaration = buildDeclaration . callStruct . typeName . SMeasure mdataPtrDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl -mdataPtrDeclaration typ = buildPtrDeclaration (callStruct (mdataName typ)) +mdataPtrDeclaration = buildPtrDeclaration . callStruct . typeName . SMeasure mdataWeight :: CExpr -> CExpr mdataWeight d = CMember d (Ident "weight") True @@ -189,39 +225,42 @@ mdataPtrSample :: CExpr -> CExpr mdataPtrSample d = CMember d (Ident "sample") False + + -------------------------------------------------------------------------------- --- | datumProd and datumSum use a store of names, which needs to match up with --- the names used when they are assigned and printed --- datumDeclaration declares struct internally --- datumStruct declares struct definitions externally +-- Datum -- +-------------------------------------------------------------------------------- +{- + In order to successfully represent Hakaru datum (Sums of Products of Hakaru + types), we must have: --- | datumName provides a unique name to identify a struct type -datumName :: Sing (a :: [[HakaruFun]]) -> String -datumName SVoid = "V" -datumName (SPlus prodD sumD) = concat ["S",datumName' prodD,datumName sumD] - where datumName' :: Sing (a :: [HakaruFun]) -> String - datumName' SDone = "U" - datumName' (SEt (SKonst x) prod') = concat ["S",tail . show $ x,datumName' prod'] - datumName' (SEt SIdent _) = error "TODO: datumName of SIdent" + > unique names for a given datum so if SVoid occurs twice in a program, C will + be using the same structure -datumNames :: [String] -datumNames = filter (\n -> not $ elem (head n) ['0'..'9']) names - where base = ['0'..'9'] ++ ['a'..'z'] - names = [[x] | x <- base] `mplus` (do n <- names - [n++[x] | x <- base]) + > C structs + > A datum may be recursive, so we will need to generate structures for all + subtypes as well. These subtypes will need to be declared before the datum + for the code to compile +-} + datumStruct :: (Sing (HData' t)) -> CExtDecl -datumStruct (SData _ typ) = CDeclExt $ datumSum typ (Ident (datumName typ)) +datumStruct dat@(SData _ typ) + = CDeclExt $ datumSum dat typ (Ident (typeName dat)) datumDeclaration :: (Sing (HData' t)) -> Ident -> CDecl -datumDeclaration (SData _ typ) = buildDeclaration (callStruct (datumName typ)) +datumDeclaration = buildDeclaration . callStruct . typeName -datumSum :: Sing (a :: [[HakaruFun]]) -> Ident -> CDecl -datumSum funs ident = - let declrs = fst $ runState (datumSum' funs) datumNames +datumSum + :: Sing (HData' t) + -> Sing (a :: [[HakaruFun]]) + -> Ident + -> CDecl +datumSum dat funs ident = + let declrs = fst $ runState (datumSum' dat funs) cNameStream union = buildDeclaration (buildUnion declrs) (Ident "sum") index = buildDeclaration CInt (Ident "index") struct = buildStruct (Just ident) $ case declrs of @@ -229,40 +268,73 @@ _ -> [index,union] in CDecl [ CTypeSpec struct ] [] -datumSum' :: Sing (a :: [[HakaruFun]]) -> State [String] [CDecl] -datumSum' SVoid = return [] -datumSum' (SPlus prod rest) = +datumSum' + :: Sing (HData' t) + -> Sing (a :: [[HakaruFun]]) + -> State [String] [CDecl] +datumSum' _ SVoid = return [] +datumSum' dat (SPlus prod rest) = do (name:names) <- get put names let ident = Ident name - mdecl = datumProd prod ident - rest' <- datumSum' rest + mdecl = datumProd dat prod ident + rest' <- datumSum' dat rest case mdecl of Nothing -> return rest' Just d -> return $ [d] ++ rest' - -datumProd :: Sing (a :: [HakaruFun]) -> Ident -> Maybe CDecl -datumProd SDone _ = Nothing -datumProd funs ident = - let declrs = fst $ runState (datumProd' funs) datumNames +datumProd + :: Sing (HData' t) + -> Sing (a :: [HakaruFun]) + -> Ident + -> Maybe CDecl +datumProd _ SDone _ = Nothing +datumProd dat funs ident = + let declrs = fst $ runState (datumProd' dat funs) cNameStream in Just $ buildDeclaration (buildStruct Nothing $ declrs) ident -- datumProd uses a store of names, which needs to match up with the names used -- when they are assigned as well as printed -datumProd' :: Sing (a :: [HakaruFun]) -> State [String] [CDecl] -datumProd' SDone = return [] -datumProd' (SEt (SKonst t) rest) = +datumProd' + :: Sing (HData' t) + -> Sing (a :: [HakaruFun]) + -> State [String] [CDecl] +datumProd' _ SDone = return [] +datumProd' dat (SEt x ps) = + do x' <- datumPrim dat x + ps' <- datumProd' dat ps + return $ x' ++ ps' + +-- We need to pass HData in case it is some recursive type +datumPrim + :: Sing (HData' t) + -> Sing (a :: HakaruFun) + -> State [String] [CDecl] +datumPrim dat prim = do (name:names) <- get put names let ident = Ident name - decl = typeDeclaration t ident - rest' <- datumProd' rest - return $ [decl] ++ rest' -datumProd' (SEt SIdent _) = error "TODO: datumProd' SIdent" + decl = case prim of + SIdent -> datumDeclaration dat ident + (SKonst k) -> typeDeclaration k ident + return [decl] ----------------------------------------------------------------- +-- index into pair +datumFst :: CExpr -> CExpr +datumFst x = x ... "sum" ... "a" ... "a" +datumSnd :: CExpr -> CExpr +datumSnd x = x ... "sum" ... "a" ... "b" + +-------------------------------------------------------------------------------- +-- Functions -- +-------------------------------------------------------------------------------- +{- + This still needs some work. Currently, we use the CodeGenMonad to give us + a list of local declarations and statements to be used in a function. Then + build a function from that. +-} + functionDef :: Sing (a :: Hakaru) -> Ident @@ -272,39 +344,52 @@ -> CFunDef functionDef typ ident argDecls internalDecls stmts = CFunDef (fmap CTypeSpec $ buildType typ) - (CDeclr Nothing [ CDDeclrIdent ident ]) + (CDeclr Nothing (CDDeclrIdent ident)) argDecls - (CCompound ((fmap CBlockDecl internalDecls) ++ (fmap CBlockStat stmts))) + (CCompound ((fmap CBlockDecl internalDecls) + ++ (fmap CBlockStat stmts))) ----------------------------------------------------------------- +-------------- +-- Closures -- +-------------- + +closureDeclaration + :: (Sing (a :: Hakaru)) + -> Ident + -> CDecl +closureDeclaration = buildDeclaration . callStruct . typeName + + + +-------------------------------------------------------------------------------- -- | buildType function do the work of describing how the Hakaru -- type will be stored in memory. Arrays needed their own -- declaration function for their arity + buildType :: Sing (a :: Hakaru) -> [CTypeSpec] -buildType SInt = [CInt] -buildType SNat = [CUnsigned, CInt] -buildType SProb = [CDouble] -buildType SReal = [CDouble] -buildType (SMeasure x) = [callStruct . mdataName $ x] -buildType (SArray t) = [callStruct $ arrayName t] -buildType (SFun _ x) = buildType $ x -- build type the function returns -buildType (SData _ t) = [callStruct $ datumName t] +buildType SInt = [CInt] +buildType SNat = [CUnsigned, CInt] +buildType SProb = [CDouble] +buildType SReal = [CDouble] +buildType (SMeasure x) = [callStruct . typeName . SMeasure $ x] +buildType (SArray t) = [callStruct . typeName . SArray $ t] +buildType (SFun _ x) = buildType $ x -- build type the function returns +buildType d@(SData _ _) = [callStruct . typeName $ d] -- these mk...Decl functions are used in coersions -mkDecl :: [CTypeSpec] -> CDecl -mkDecl t = CDecl (fmap CTypeSpec t) [] +castTo :: [CTypeSpec] -> CExpr -> CExpr +castTo t = CCast (CTypeName t False) -mkPtrDecl :: [CTypeSpec] -> CDecl -mkPtrDecl t = CDecl (fmap CTypeSpec t) - [( CDeclr (Just $ CPtrDeclr []) [] - , Nothing )] +castToPtrOf :: [CTypeSpec] -> CExpr -> CExpr +castToPtrOf t = CCast (CTypeName t True) buildStruct :: Maybe Ident -> [CDecl] -> CTypeSpec buildStruct mi decls = CSUType (CSUSpec CStructTag mi decls) --- | callStruct will give the type spec calling a struct we have already declared externally +-- | callStruct will give the type spec calling a struct we have already +-- declared externally callStruct :: String -> CTypeSpec callStruct name = CSUType (CSUSpec CStructTag (Just (Ident name)) []) @@ -314,31 +399,8 @@ CSUType (CSUSpec CUnionTag Nothing decls) -intDecl, natDecl, doubleDecl :: CDecl -intDecl = mkDecl [CInt] -natDecl = CDecl [CTypeSpec CUnsigned - ,CTypeSpec CInt] - [( CDeclr (Just $ CPtrDeclr []) [], Nothing )] -doubleDecl = mkDecl [CDouble] - -intPtr, natPtr, doublePtr :: CDecl -intPtr = mkPtrDecl [CInt] -natPtr = CDecl [CTypeSpec CUnsigned, CTypeSpec CInt] [] -doublePtr = mkPtrDecl [CDouble] - -boolTyp :: CDecl -boolTyp = - CDecl [CTypeSpec - (CSUType - (CSUSpec CStructTag - (Just (Ident "bool")) - [buildDeclaration CInt (Ident "index")]))] - [] - - - binaryOp :: NaryOp a -> CExpr -> CExpr -> CExpr -binaryOp (Sum HSemiring_Prob) a b = CBinary CAddOp (exp a) (exp b) +binaryOp (Sum HSemiring_Prob) a b = CBinary CAddOp (expE a) (expE b) binaryOp (Prod HSemiring_Prob) a b = CBinary CAddOp a b binaryOp (Sum _) a b = CBinary CAddOp a b binaryOp (Prod _) a b = CBinary CMulOp a b diff --git a/haskell/Language/Hakaru/CodeGen/Wrapper.hs b/haskell/Language/Hakaru/CodeGen/Wrapper.hs --- a/haskell/Language/Hakaru/CodeGen/Wrapper.hs +++ b/haskell/Language/Hakaru/CodeGen/Wrapper.hs @@ -6,7 +6,8 @@ GADTs, KindSignatures, RankNTypes, - ScopedTypeVariables #-} + ScopedTypeVariables, + TypeOperators #-} ---------------------------------------------------------------- -- 2016.06.23 @@ -39,8 +40,8 @@ import Language.Hakaru.CodeGen.Flatten import Language.Hakaru.CodeGen.Types import Language.Hakaru.CodeGen.AST +import Language.Hakaru.CodeGen.Libs import Language.Hakaru.Types.DataKind (Hakaru(..)) - import Control.Monad.State.Strict import Prelude as P hiding (unlines,exp) @@ -59,29 +60,24 @@ -> Maybe String -- ^ Maybe an output name -> PrintConfig -- ^ show weights? -> CodeGen () -wrapProgram tast@(TypedAST typ _) mn pc = +wrapProgram tast@(TypedAST typ _) mn pconfig = do sequence_ . fmap (extDeclare . CPPExt) . header $ typ - baseCG - return () - where baseCG = case (tast,mn) of - ( TypedAST (SFun _ _) abt, Just name ) -> - do reserveName name - flattenTopLambda abt $ Ident name - - ( TypedAST (SFun _ _) abt, Nothing ) -> - genIdent' "fn" >>= \name -> - flattenTopLambda abt name - + cg <- get + when (managedMem cg) $ extDeclare . CPPExt $ gcHeader + when (sharedMem cg) $ extDeclare . CPPExt $ openMpHeader + case (tast,mn) of + ( TypedAST (SFun _ _) abt, Just name ) -> + flattenTopLambda abt =<< reserveIdent name - ( TypedAST _ _, Just _ ) -> undefined - -- do reserveName name - -- defineFunction typ' - -- (Ident name) - -- [] - -- (putStat . CReturn . Just =<< flattenABT abt) + ( TypedAST typ' abt, Just name ) -> + -- still buggy for measures + do mfId <- reserveIdent name + funCG (head . buildType $ typ') mfId [] $ + do outE <- flattenWithName' abt "out" + putStat . CReturn . Just $ outE - ( TypedAST typ' abt, Nothing ) -> - mainFunction pc typ' abt + ( TypedAST typ' abt, Nothing ) -> + mainFunction pconfig typ' abt @@ -109,70 +105,219 @@ -> CodeGen () -- when measure, compile to a sampler -mainFunction pc typ@(SMeasure t) abt = - let ident = Ident "measure" - funId = Ident "main" - mdataId = Ident "mdata" - in do reserveName "measure" - reserveName "mdata" - reserveName "main" +mainFunction pconfig typ@(SMeasure _) abt = + do mfId <- reserveIdent "measure" + mdataId <- reserveIdent "mdata" + mainId <- reserveIdent "main" + argVId <- reserveIdent "argv" + argCId <- reserveIdent "argc" + let (argCE:argVE:[]) = fmap CVar [argCId,argVId] + extDeclareTypes typ - extDeclare . mdataStruct $ t + -- defined a measure function that returns mdata + funCG (head . buildType $ typ) mfId [] $ + (putStat . CReturn . Just) =<< flattenWithName' abt "samp" - -- defined a measure function that returns mdata - funCG CVoid ident [mdataPtrDeclaration t mdataId] $ - do flattenABT abt (CVar mdataId) - putStat (CReturn Nothing) + funCG CInt mainId mainArgs $ + do isManagedMem <- managedMem <$> get + when isManagedMem (putExprStat gcInit) - -- -- need to set seed? - -- -- srand(time(NULL)); + nSamples <- parseNumSamples argCE argVE + seedE <- parseSeed argCE argVE - printf pc typ (CVar ident) - putStat . CReturn . Just $ intE 0 + putExprStat $ mkCallE "srand" [seedE] - !cg <- get - extDeclare . CFunDefExt $ functionDef SInt - funId - [] - (P.reverse $ declarations cg) - (P.reverse $ statements cg) - -- where isSArray (SArray _) = True - -- isSArray _ = False - -- mkArrayStruct :: Sing (a :: Hakaru) -> CExtDecl - -- mkArrayStruct (SArray t) = arrayStruct t - -- mkArrayStruct _ = error "Not Array" - -- getArrayType :: Sing (b :: Hakaru) -> [CTypeSpec] - -- getArrayType (SArray t) = case buildType t of - -- [] -> error "wrapper: this shouldn't happen" - -- t -> t - -- getArrayType _ = error "Not Array" - -- getPlateArity :: ABT Term abt => Term abt a -> abt '[] 'HNat - -- getPlateArity (Plate :$ arity :* _ :* End) = arity - -- getPlateArity _ = error "mainFunction not a plate" + putStat $ opComment "Run Hakaru Sampler" + printCG pconfig typ (CCall (CVar mfId) []) (Just nSamples) + putStat . CReturn . Just $ intE 0 +mainFunction pconfig typ@(SFun _ _) abt = + coalesceLambda abt $ \vars abt' -> + do resId <- reserveIdent "result" + mainId <- reserveIdent "main" + argVId <- reserveIdent "argv" + argCId <- reserveIdent "argc" + funId <- genIdent' "fn" + flattenTopLambda abt funId + let (resE:funE:argCE:argVE:[]) = fmap CVar [resId,funId,argCId,argVId] + typ' = typeOf abt' + + funCG CInt mainId mainArgs $ + do isManagedMem <- managedMem <$> get + when isManagedMem (putExprStat gcInit) + declare typ' resId + + mns <- maybeNumSamples argCE argVE typ' + case typ' of + SMeasure _ -> do seedE <- parseSeed argCE argVE + putExprStat $ mkCallE "srand" [seedE] + _ -> return () + + withLambdaDepth' 0 abt $ \d -> + let argErr 0 = "" + argErr n = (argErr (pred n)) ++ " " in + ifCG (argCE .<. (intE (d+1))) + (do putExprStat $ printfE + [ stringE $ "Usage: %s " ++ argErr d ++ "\n" + , (index argVE (intE 0)) ] + putExprStat $ mkCallE "abort" [ ]) + (return ()) + + putStat $ opComment "Parse Args" + argEs <- foldLambdaWithIndex 1 abt $ \i (Variable _ _ t) -> + do argE <- localVar' t "arg" + parseCG t (index argVE (intE i)) argE + return argE + + case typ' of + SMeasure _ -> do putStat $ opComment "Run Hakaru Sampler" + printCG pconfig typ' (CCall funE argEs) mns + + _ -> do putStat $ opComment "Run Hakaru Program" + putExprStat $ resE .=. (CCall funE argEs) + putStat $ opComment "Print Result" + printCG pconfig typ' resE mns + + putStat . CReturn . Just $ intE 0 + + + where withLambdaDepth' + :: ABT Term abt + => Integer + -> abt '[] a + -> ( forall (x :: Hakaru) + . Integer + -> CodeGen ()) + -> CodeGen () + withLambdaDepth' n abt_ k = + caseVarSyn abt_ + (const (k n)) + (\term -> + case term of + (Lam_ :$ body :* End) -> + caseBind body $ \_ abt_' -> + withLambdaDepth' (succ n) abt_' k + _ -> k n) + + maybeNumSamples + :: CExpr -> CExpr -> Sing (a :: Hakaru) -> CodeGen (Maybe CExpr) + maybeNumSamples c v (SMeasure _) = Just <$> parseNumSamples c v + maybeNumSamples _ _ _ = return Nothing + -- just a computation -mainFunction pc typ abt = - let resId = Ident "result" - resE = CVar resId - funId = Ident "main" - in do reserveName "result" - reserveName "main" +mainFunction pconfig typ abt = + do resId <- reserveIdent "result" + mainId <- reserveIdent "main" + let resE = CVar resId - declare typ resId - flattenABT abt resE + funCG CInt mainId [] $ + do declare typ resId - printf pc typ resE - putStat . CReturn . Just $ intE 0 + isManagedMem <- managedMem <$> get + when isManagedMem (putExprStat gcInit) - cg <- get - extDeclare . CFunDefExt $ functionDef SInt - funId - [] - (P.reverse $ declarations cg) - (P.reverse $ statements cg) + flattenABT abt resE + printCG pconfig typ resE Nothing + putStat . CReturn . Just $ intE 0 +mainArgs :: [CDecl] +mainArgs = [ CDecl [CTypeSpec CInt] + [(CDeclr Nothing (CDDeclrIdent (Ident "argc")), Nothing)] + , CDecl [CTypeSpec CChar] + [(CDeclr (Just (CPtrDeclr [])) + (CDDeclrArr (CDDeclrIdent (Ident "argv")) Nothing) + , Nothing)] + ] +{- the number of samples is set to -1 by default -} +parseNumSamples :: CExpr -> CExpr -> CodeGen CExpr +parseNumSamples argc argv = + do itE <- localVar SNat + outE <- localVar SNat + putStat $ opComment "Num Samples?" + putExprStat $ outE .=. (intE (-1)) + forCG (itE .=. (intE 1)) + (itE .<. argc) + (CUnary CPostIncOp itE) + (ifCG ((index (index argv itE) (intE 0) .==. (charE '-')) .&&. + (index (index argv itE) (intE 1) .==. (charE 'n'))) + (putExprStat $ + sscanfE [index argv itE,stringE "-n%d",address outE]) + (return ())) + putStat $ opComment "End Num Samples?" + return outE + +{- the randome seed is set to time(NULL) by default -} +parseSeed :: CExpr -> CExpr -> CodeGen CExpr +parseSeed argc argv = + do itE <- localVar SNat + outE <- localVar SNat + putStat $ opComment "Random Seed?" + putExprStat $ outE .=. (mkCallE "time" [ CVar . Ident $ "NULL"]) + forCG (itE .=. (intE 1)) + (itE .<. argc) + (CUnary CPostIncOp itE) + (ifCG ((index (index argv itE) (intE 0) .==. (charE '-')) .&&. + (index (index argv itE) (intE 1) .==. (charE 's'))) + (putExprStat $ + sscanfE [index argv itE,stringE "-s%d",address outE]) + (return ())) + putStat $ opComment "End Random Seed?" + return outE + + -------------------------------------------------------------------------------- +-- Parsing Values -- +-------------------------------------------------------------------------------- + +parseCG :: Sing (a :: Hakaru) -> CExpr -> CExpr -> CodeGen CExpr +parseCG (SArray t) from to = + do fpId <- genIdent' "fp" + buffId <- genIdent' "buff" + declare' $ CDecl [CTypeSpec fileT] + [(CDeclr (Just (CPtrDeclr [])) + (CDDeclrIdent fpId) + , Nothing)] + declare' $ CDecl [CTypeSpec CChar] + [(CDeclr Nothing + (CDDeclrArr (CDDeclrIdent buffId) (Just (intE 1024))) + , Nothing)] + let fpE = CVar fpId + buffE = CVar buffId + putExprStat $ fpE .=. (fopenE from (stringE "r")) + itE <- localVar SNat + putExprStat $ itE .=. (intE 0) + whileCG (fgetsE buffE (intE 1024) fpE .!=. nullE) + (putExprStat $ CUnary CPostIncOp itE) + putExprStat $ arraySize to .=. itE + putMallocStat (arrayData to) itE t + putExprStat $ itE .=. (intE 0) + putExprStat $ rewindE fpE + whileCG (fgetsE buffE (intE 1024) fpE .!=. nullE) + (do checkE <- parseCG t buffE (index (arrayData to) itE) + ifCG (checkE .==. (intE 1)) + (putExprStat $ CUnary CPostIncOp itE) + (putExprStat $ CUnary CPostDecOp (arraySize to))) + putExprStat $ fcloseE fpE + localVar SNat + +parseCG t from to = + do checkE <- localVar SNat + putExprStat $ checkE .=. sscanfE [from,stringE . parseFormat $ t,address to] + case t of + SProb -> putExprStat $ to .=. logE to + _ -> return () + return checkE + +parseFormat :: Sing (a :: Hakaru) -> String +parseFormat SInt = "%d" +parseFormat SNat = "%u" +parseFormat SReal = "%lf" +parseFormat SProb = "%lf" +parseFormat t = error $ "parseCG{" ++ show t ++ "}: no available parsing form" + + +-------------------------------------------------------------------------------- -- Printing Values -- -------------------------------------------------------------------------------- {- @@ -183,92 +328,73 @@ -} data PrintConfig - = PrintConfig { showWeights :: Bool + = PrintConfig { noWeights :: Bool , showProbInLog :: Bool } deriving Show -printf +printCG :: PrintConfig -> Sing (a :: Hakaru) -- ^ Hakaru type to be printed -> CExpr -- ^ CExpr representing value + -> Maybe CExpr -- ^ If measure type, expr for num samples -> CodeGen () - -printf pc mt@(SMeasure t) sampleFunc = - case t of - _ -> do mId <- genIdent' "m" - declare mt mId - let mE = CVar mId - getSampleS = CExpr . Just $ CCall sampleFunc [address mE] - printSampleE = CExpr . Just - $ CCall (CVar . Ident $ "printf") - $ [ stringE $ printfText pc mt "\n"] - ++ (if showWeights pc - then [ if showProbInLog pc - then mdataWeight mE - else exp $ mdataWeight mE ] - else []) - ++ [ case t of - SProb -> if showProbInLog pc - then mdataSample mE - else exp $ mdataSample mE - _ -> mdataSample mE ] - wrapSampleFunc = CCompound $ [CBlockStat getSampleS - ,CBlockStat $ CIf ((exp $ mdataWeight mE) .>. (floatE 0)) printSampleE Nothing] - putStat $ CWhile (intE 1) wrapSampleFunc False - - -printf pc (SArray t) arg = - do iterId <- genIdent' "it" - declare SInt iterId - let iter = CVar iterId - result = arg - dataPtr = CMember result (Ident "data") True - sizeVar = CMember result (Ident "size") True - cond = iter .<. sizeVar - inc = CUnary CPostIncOp iter - currInd = indirect (dataPtr .+. iter) - loopBody = putExprStat $ CCall (CVar . Ident $ "printf") - [ stringE $ printfText pc t " " - , currInd ] - +printCG pconfig mtyp@(SMeasure typ) sampleFunc (Just numSamples) = + do mE <- localVar' mtyp "m" + itE <- localVar SNat + putExprStat $ itE .=. numSamples + whileCG (itE .!=. (intE 0)) $ + do putExprStat $ mE .=. sampleFunc + case noWeights pconfig of + True -> printCG pconfig typ (mdataSample mE) Nothing + False -> do putExprStat $ + printfE [ stringE (printFormat pconfig SProb "\t") + , expE (mdataWeight mE) ] + printCG pconfig typ (mdataSample mE) Nothing + ifCG (numSamples .>=. (intE 0)) + (putExprStat $ CUnary CPostDecOp itE) + (return ()) +printCG pconfig (SArray typ) arg Nothing = + do itE <- localVar' SNat "it" putString "[ " - mkSequential -- cant print arrays in parallel - forCG (iter .=. (intE 0)) cond inc loopBody + mkSequential + forCG (itE .=. (intE 0)) + (itE .<. (arraySize arg)) + (CUnary CPostIncOp itE) + (putExprStat + $ printfE [ stringE $ printFormat pconfig typ " " + , index (arrayData arg) itE ]) putString "]\n" - where putString s = putExprStat $ CCall (CVar . Ident $ "printf") - [stringE s] + where putString s = putExprStat $ printfE [stringE s] -printf pc SProb arg = - putExprStat $ CCall (CVar . Ident $ "printf") - [ stringE $ printfText pc SProb "\n" - , if showProbInLog pc +printCG pconfig SProb arg Nothing = + putExprStat $ printfE + [ stringE $ printFormat pconfig SProb "\n" + , if showProbInLog pconfig then arg - else exp arg ] - -printf pc typ arg = - putExprStat $ CCall (CVar . Ident $ "printf") - [ stringE $ printfText pc typ "\n" - , arg ] - + else expE arg ] +printCG pconfig typ arg Nothing = + putExprStat $ printfE + [ stringE $ printFormat pconfig typ "\n" + , arg ] -printfText :: PrintConfig -> Sing (a :: Hakaru) -> (String -> String) -printfText _ SInt = \s -> "%d" ++ s -printfText _ SNat = \s -> "%d" ++ s -printfText c SProb = \s -> if showProbInLog c +printFormat :: PrintConfig -> Sing (a :: Hakaru) -> (String -> String) +printFormat _ SInt = \s -> "%d" ++ s +printFormat _ SNat = \s -> "%d" ++ s +printFormat c SProb = \s -> if showProbInLog c then "exp(%.15f)" ++ s else "%.15f" ++ s -printfText _ SReal = \s -> "%.17f" ++ s -printfText c (SMeasure t) = if showWeights c - then \s -> if showProbInLog c - then "exp(%.15f) " ++ printfText c t s - else "%.15f " ++ printfText c t s - else printfText c t -printfText c (SArray t) = printfText c t -printfText _ (SFun _ _) = id -printfText _ (SData _ _) = \s -> "TODO: printft datum" ++ s +printFormat _ SReal = \s -> "%.15f" ++ s +printFormat c (SMeasure t) = if noWeights c + then printFormat c t + else \s -> if showProbInLog c + then "exp(%.15f) " ++ printFormat c t s + else "%.15f " ++ printFormat c t s +printFormat c (SArray t) = printFormat c t +printFormat _ (SFun _ _) = id +printFormat _ (SData _ _) = \s -> "TODO: printft datum" ++ s -------------------------------------------------------------------------------- @@ -291,59 +417,57 @@ -> CodeGen () flattenTopLambda abt name = coalesceLambda abt $ \vars abt' -> - let varMs = foldMap11 (\v -> [mkVarDecl v =<< createIdent v]) vars + let varMs = foldMap11 (\v -> [mkVarDecl v =<< createIdent' "param" v]) vars typ = typeOf abt' in do argDecls <- sequence varMs - cg <- get - case typ of - -- SMeasure _ -> error "flattenTopLambda: for Measures" - -- SMeasure _ -> do let m = putStat . CReturn . Just =<< flattenABT abt' - -- (_,cg') = runState m $ cg { statements = [] - -- , declarations = [] } - -- put $ cg' { statements = statements cg - -- , declarations = declarations cg } - -- extDeclare . CFunDefExt - -- $ functionDef typ name - -- argDecls - -- (P.reverse $ declarations cg') - -- (P.reverse $ statements cg') - _ -> do let m = do outId <- genIdent' "out" - declare (typeOf abt') outId - let outE = CVar outId - flattenABT abt' outE - putStat . CReturn . Just $ outE - (_,cg') = runState m $ cg { statements = [] - , declarations = [] } - put $ cg' { statements = statements cg - , declarations = declarations cg } - extDeclare . CFunDefExt - $ functionDef typ name - argDecls - (P.reverse $ declarations cg') - (P.reverse $ statements cg') - -- do at top level - where coalesceLambda - :: ABT Term abt - => abt '[] a - -> ( forall (ys :: [Hakaru]) b - . List1 Variable ys -> abt '[] b -> r) - -> r - coalesceLambda abt_ k = - caseVarSyn abt_ (const (k Nil1 abt_)) $ \term -> - case term of - (Lam_ :$ body :* End) -> - caseBind body $ \v body' -> - coalesceLambda body' $ \vars body'' -> k (Cons1 v vars) body'' - _ -> k Nil1 abt_ - + funCG (head . buildType $ typ) name argDecls $ + (putStat . CReturn . Just) =<< flattenWithName' abt' "out" - mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CodeGen CDecl + -- do at top level + where mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CodeGen CDecl mkVarDecl (Variable _ _ SInt) = return . typeDeclaration SInt mkVarDecl (Variable _ _ SNat) = return . typeDeclaration SNat mkVarDecl (Variable _ _ SProb) = return . typeDeclaration SProb mkVarDecl (Variable _ _ SReal) = return . typeDeclaration SReal - mkVarDecl (Variable _ _ (SArray t)) = \i -> do extDeclare $ arrayStruct t - return $ arrayDeclaration t i - mkVarDecl (Variable _ _ d@(SData _ _)) = \i -> do extDeclare $ datumStruct d - return $ datumDeclaration d i + mkVarDecl (Variable _ _ (SArray t)) = \i -> + extDeclareTypes (SArray t) >> (return $ arrayDeclaration t i) + + mkVarDecl (Variable _ _ d@(SData _ _)) = \i -> + extDeclareTypes d >> (return $ datumDeclaration d i) + mkVarDecl v = error $ "flattenSCon.Lam_.mkVarDecl cannot handle vars of type " ++ show v + +coalesceLambda + :: ABT Term abt + => abt '[] a + -> ( forall (ys :: [Hakaru]) b + . List1 Variable ys -> abt '[] b -> r) + -> r +coalesceLambda abt_ k = + caseVarSyn abt_ (const (k Nil1 abt_)) $ \term -> + case term of + (Lam_ :$ body :* End) -> + caseBind body $ \v body' -> + coalesceLambda body' $ \vars body'' -> k (Cons1 v vars) body'' + _ -> k Nil1 abt_ + +foldLambdaWithIndex + :: ABT Term abt + => Integer + -> abt '[] a + -> ( forall (x :: Hakaru) + . Integer + -> Variable x + -> CodeGen CExpr) + -> CodeGen [CExpr] +foldLambdaWithIndex n abt_ k = + caseVarSyn abt_ + (const (return [])) + (\term -> + case term of + (Lam_ :$ body :* End) -> + caseBind body $ \v abt_' -> + (do x <- k n v + xs <- foldLambdaWithIndex (succ n) abt_' k + return (x:xs)) + _ -> return []) diff --git a/haskell/Language/Hakaru/Command.hs b/haskell/Language/Hakaru/Command.hs --- a/haskell/Language/Hakaru/Command.hs +++ b/haskell/Language/Hakaru/Command.hs @@ -1,16 +1,30 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} module Language.Hakaru.Command where import Language.Hakaru.Syntax.ABT import qualified Language.Hakaru.Syntax.AST as T +import Language.Hakaru.Parser.Import (expandImports) import Language.Hakaru.Parser.Parser hiding (style) import Language.Hakaru.Parser.SymbolResolve (resolveAST) import Language.Hakaru.Syntax.TypeCheck -import qualified Data.Text as Text -import qualified Data.Text.IO as IO +import Control.Monad.Trans.Except +import Control.Monad (when) +import qualified Data.Text as Text +import qualified Data.Text.IO as IO +import qualified Data.Text.Utf8 as U +import qualified Options.Applicative as O import Data.Vector +import System.IO (stderr) +import System.Environment (getArgs) +import Data.Monoid ((<>),mconcat) +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + type Term a = TrivialABT T.Term '[] a parseAndInfer :: Text.Text @@ -22,13 +36,60 @@ let m = inferType (resolveAST past) in runTCM m (splitLines x) LaxMode +parseAndInfer' :: Text.Text + -> IO (Either Text.Text (TypedAST (TrivialABT T.Term))) +parseAndInfer' x = + case parseHakaruWithImports x of + Left err -> return . Left $ Text.pack . show $ err + Right past -> do + past' <- runExceptT (expandImports past) + case past' of + Left err -> return . Left $ Text.pack . show $ err + Right past'' -> do + let m = inferType (resolveAST past'') + return (runTCM m (splitLines x) LaxMode) + +parseAndInferWithDebug + :: Bool + -> Text.Text + -> IO (Either Text.Text (TypedAST (TrivialABT T.Term))) +parseAndInferWithDebug debug x = + case parseHakaru x of + Left err -> return $ Left (Text.pack . show $ err) + Right past -> do + when debug $ putErrorLn $ hrule "Parsed AST" + when debug $ putErrorLn . Text.pack . show $ past + let resolved = resolveAST past + let inferred = runTCM (inferType resolved) (splitLines x) LaxMode + when debug $ putErrorLn $ hrule "Inferred AST" + when debug $ putErrorLn . Text.pack . show $ inferred + return $ inferred + where hrule s = Text.concat ["\n<=======================| " + ,s," |=======================>\n"] + putErrorLn = IO.hPutStrLn stderr + + splitLines :: Text.Text -> Maybe (Vector Text.Text) splitLines = Just . fromList . Text.lines readFromFile :: String -> IO Text.Text -readFromFile "-" = IO.getContents -readFromFile x = IO.readFile x +readFromFile "-" = U.getContents +readFromFile x = U.readFile x +simpleCommand :: (Text.Text -> IO ()) -> Text.Text -> IO () +simpleCommand k fnName = + let parser = + O.info (O.helper <*> opts) + (O.fullDesc <> O.progDesc + (mconcat["Hakaru:", Text.unpack fnName, " command"])) + opts = + O.strArgument + ( O.metavar "PROGRAM" <> + O.showDefault <> O.value "-" <> + O.help "Filepath to Hakaru program OR \"-\"" ) + + in O.execParser parser >>= readFromFile >>= k + writeToFile :: String -> (Text.Text -> IO ()) -writeToFile "-" = IO.putStrLn -writeToFile x = IO.writeFile x +writeToFile "-" = U.putStrLn +writeToFile x = U.writeFile x diff --git a/haskell/Language/Hakaru/Disintegrate.hs b/haskell/Language/Hakaru/Disintegrate.hs --- a/haskell/Language/Hakaru/Disintegrate.hs +++ b/haskell/Language/Hakaru/Disintegrate.hs @@ -82,9 +82,9 @@ -- statement. ---------------------------------------------------------------- module Language.Hakaru.Disintegrate - ( + ( lam_ -- * the Hakaru API - disintegrateWithVar + , disintegrateWithVar , disintegrate , densityWithVar , density @@ -115,8 +115,7 @@ import Data.Sequence (Seq) import qualified Data.Sequence as S import Data.Proxy (KProxy(..)) -import qualified Data.Set as Set (fromList) -import Data.Maybe (fromMaybe) +import Data.Maybe (fromMaybe, fromJust) import Language.Hakaru.Syntax.IClasses import Data.Number.Natural @@ -130,7 +129,7 @@ import Language.Hakaru.Syntax.DatumCase (DatumEvaluator, MatchResult(..), matchBranches) import Language.Hakaru.Syntax.ABT import Language.Hakaru.Evaluation.Types -import Language.Hakaru.Evaluation.Lazy hiding (evaluate,update) +import Language.Hakaru.Evaluation.Lazy import Language.Hakaru.Evaluation.DisintegrationMonad import qualified Language.Hakaru.Syntax.Prelude as P import qualified Language.Hakaru.Expect as E @@ -185,31 +184,6 @@ ++ show (pretty_Statements ss PP.$+$ PP.sep(prettyPrec_ 11 ab)) ++ "\n") $ return () #endif - -- BUG: Why does 'testDisintegrate1a' return no solutions? - -- - -- In older code (up to git#38889a5): It's because 'emitUnpair' - -- isn't quite smart enough. When the @ab@ expression is a - -- 'Neutral' case expression, we need to go underneath the - -- case expression and call 'constrainValue' on each branch. - -- Instead, what we currently do is emit an @unpair@ case - -- statement with the scrutinee being the 'Neutral' case - -- expression, and then just return the pair of variables - -- bound by the emitted @unpair@; but, of course, - -- 'constrainValue' can't do anything with those variables - -- (since they appear to be free, given as they've already - -- been emitted). Another way to think about what it is we - -- need to do to correct this is that we need to perform - -- the case-of-case transformation (where one of the cases - -- is the 'Neutral' one, and the other is the @unpair@). - -- - -- In newer code (git#e8a0c66 and later): When we call 'perform' - -- on an 'SBind' statement we emit some code and update the - -- binding to become an 'SLet' of some local variable to - -- the emitted variable. Later on when we call 'constrainVariable' - -- on the local variable, we will look that 'SLet' statement - -- up; and then when we call 'constrainVariable' on the - -- emitted variable, things will @bot@ because we cannot - -- constrain free variables in general. (a,b) <- emitUnpair ab #ifdef __TRACE_DISINTEGRATE__ trace ("-- disintegrate: finished emitUnpair: " @@ -218,10 +192,10 @@ constrainValue (var x) a #ifdef __TRACE_DISINTEGRATE__ ss <- getStatements - locs <- getLocs + extras <- getExtras traceM ("-- disintegrate: finished constrainValue\n" ++ show (pretty_Statements ss) ++ "\n" - ++ show (prettyLocs locs) + ++ show (prettyExtras extras) ) #endif return b @@ -296,623 +270,19 @@ ---------------------------------------------------------------- ---------------------------------------------------------------- -firstM :: Functor f => (a -> f b) -> (a,c) -> f (b,c) -firstM f (x,y) = (\z -> (z, y)) <$> f x - -- N.B., forward disintegration is not identical to partial evaluation, -- as noted at the top of the file. For correctness we need to -- ensure the result is emissible (i.e., has no heap-bound variables). -- More specifically, we need to ensure emissibility in the places -- where we call 'emitMBind' evaluate_ :: (ABT Term abt) => TermEvaluator abt (Dis abt) -evaluate_ = evaluate perform evaluateCase - - --- Copying `evaluate` and `update` from LH.Evaluation.Lazy for now (2016-06-28) --- Beginning of copied code ------------------------------------------------- - -evaluate - :: forall abt m p - . (ABT Term abt) - => MeasureEvaluator abt (Dis abt) - -> (TermEvaluator abt (Dis abt) -> CaseEvaluator abt (Dis abt)) - -> TermEvaluator abt (Dis abt) -{-# INLINE evaluate #-} -evaluate perform evaluateCase = goEvaluate - where - evaluateCase_ :: CaseEvaluator abt (Dis abt) - evaluateCase_ = evaluateCase goEvaluate - - goEvaluate :: TermEvaluator abt (Dis abt) - goEvaluate e0 = -#ifdef __TRACE_DISINTEGRATE__ - getIndices >>= \inds -> - trace ("-- goEvaluate: " ++ show (pretty e0) - ++ "at " ++ show (ppInds inds)) $ -#endif - caseVarSyn e0 (update perform goEvaluate) $ \t -> - case t of - -- Things which are already WHNFs - Literal_ v -> return . Head_ $ WLiteral v - Datum_ d -> return . Head_ $ WDatum d - Empty_ typ -> return . Head_ $ WEmpty typ - Array_ e1 e2 -> return . Head_ $ WArray e1 e2 - Lam_ :$ e1 :* End -> return . Head_ $ WLam e1 - Dirac :$ e1 :* End -> return . Head_ $ WDirac e1 - MBind :$ e1 :* e2 :* End -> return . Head_ $ WMBind e1 e2 - Plate :$ e1 :* e2 :* End -> return . Head_ $ WPlate e1 e2 - MeasureOp_ o :$ es -> return . Head_ $ WMeasureOp o es - Superpose_ pes -> return . Head_ $ WSuperpose pes - Reject_ typ -> return . Head_ $ WReject typ - -- We don't bother evaluating these, even though we could... - Integrate :$ e1 :* e2 :* e3 :* End -> - return . Head_ $ WIntegrate e1 e2 e3 - Summate h1 h2 :$ e1 :* e2 :* e3 :* End -> - return . Neutral $ syn t - --return . Head_ $ WSummate e1 e2 e3 - - - -- Everything else needs some evaluation - - App_ :$ e1 :* e2 :* End -> do - w1 <- goEvaluate e1 - case w1 of - Neutral e1' -> return . Neutral $ P.app e1' e2 - Head_ v1 -> evaluateApp v1 - where - evaluateApp (WLam f) = - -- call-by-name: - caseBind f $ \x f' -> do - i <- getIndices - push (SLet x (Thunk e2) i) f' goEvaluate - evaluateApp _ = error "evaluate{App_}: the impossible happened" - - Let_ :$ e1 :* e2 :* End -> do - i <- getIndices - caseBind e2 $ \x e2' -> - push (SLet x (Thunk e1) i) e2' goEvaluate - - CoerceTo_ c :$ e1 :* End -> C.coerceTo c <$> goEvaluate e1 - UnsafeFrom_ c :$ e1 :* End -> C.coerceFrom c <$> goEvaluate e1 - - -- TODO: will maybe clean up the code to map 'evaluate' over @es@ before calling the evaluateFooOp helpers? - NaryOp_ o es -> evaluateNaryOp goEvaluate o es - ArrayOp_ o :$ es -> evaluateArrayOp goEvaluate o es - PrimOp_ o :$ es -> evaluatePrimOp goEvaluate o es - - -- BUG: avoid the chance of looping in case 'E.expect' residualizes! - -- TODO: use 'evaluate' in 'E.expect' for the evaluation of @e1@ - Expect :$ e1 :* e2 :* End -> - error "TODO: evaluate{Expect}: unclear how to handle this without cyclic dependencies" - {- - -- BUG: can't call E.expect because of cyclic dependency - goEvaluate . E.expect e1 $ \e3 -> - syn (Let_ :$ e3 :* e2 :* End) - -} - - Case_ e bs -> evaluateCase_ e bs - - _ :$ _ -> error "evaluate: the impossible happened" - -evaluateNaryOp - :: (ABT Term abt) - => TermEvaluator abt (Dis abt) - -> NaryOp a - -> Seq (abt '[] a) - -> Dis abt (Whnf abt a) -evaluateNaryOp evaluate_ = \o es -> mainLoop o (evalOp o) S.empty es - where - -- TODO: there's got to be a more efficient way to do this... - mainLoop o op ws es = - case S.viewl es of - S.EmptyL -> return $ - case S.viewl ws of - S.EmptyL -> identityElement o -- Avoid empty naryOps - w S.:< ws' - | S.null ws' -> w -- Avoid singleton naryOps - | otherwise -> - Neutral . syn . NaryOp_ o $ fmap fromWhnf ws - e S.:< es' -> do - w <- evaluate_ e - case matchNaryOp o w of - Nothing -> mainLoop o op (snocLoop op ws w) es' - Just es2 -> mainLoop o op ws (es2 S.>< es') - - snocLoop - :: (ABT syn abt) - => (Head abt a -> Head abt a -> Head abt a) - -> Seq (Whnf abt a) - -> Whnf abt a - -> Seq (Whnf abt a) - snocLoop op ws w1 = - -- TODO: immediately return @ws@ if @w1 == identityElement o@ (whenever identityElement is defined) - case S.viewr ws of - S.EmptyR -> S.singleton w1 - ws' S.:> w2 -> - case (w1,w2) of - (Head_ v1, Head_ v2) -> snocLoop op ws' (Head_ (op v1 v2)) - _ -> ws S.|> w1 - - matchNaryOp - :: (ABT Term abt) - => NaryOp a - -> Whnf abt a - -> Maybe (Seq (abt '[] a)) - matchNaryOp o w = - case w of - Head_ _ -> Nothing - Neutral e -> - caseVarSyn e (const Nothing) $ \t -> - case t of - NaryOp_ o' es | o' == o -> Just es - _ -> Nothing - - -- TODO: move this off to Prelude.hs or somewhere... - identityElement :: (ABT Term abt) => NaryOp a -> Whnf abt a - identityElement o = - case o of - And -> Head_ (WDatum dTrue) - Or -> Head_ (WDatum dFalse) - Xor -> Head_ (WDatum dFalse) - Iff -> Head_ (WDatum dTrue) - Min _ -> Neutral (syn (NaryOp_ o S.empty)) -- no identity in general (but we could do it by cases...) - Max _ -> Neutral (syn (NaryOp_ o S.empty)) -- no identity in general (but we could do it by cases...) - -- TODO: figure out how to reuse 'P.zero_' and 'P.one_' here; requires converting thr @(syn . Literal_)@ into @(Head_ . WLiteral)@. Maybe we should change 'P.zero_' and 'P.one_' so they just return the 'Literal' itself rather than the @abt@? - Sum HSemiring_Nat -> Head_ (WLiteral (LNat 0)) - Sum HSemiring_Int -> Head_ (WLiteral (LInt 0)) - Sum HSemiring_Prob -> Head_ (WLiteral (LProb 0)) - Sum HSemiring_Real -> Head_ (WLiteral (LReal 0)) - Prod HSemiring_Nat -> Head_ (WLiteral (LNat 1)) - Prod HSemiring_Int -> Head_ (WLiteral (LInt 1)) - Prod HSemiring_Prob -> Head_ (WLiteral (LProb 1)) - Prod HSemiring_Real -> Head_ (WLiteral (LReal 1)) - - -- | The evaluation interpretation of each NaryOp - evalOp - :: (ABT Term abt) - => NaryOp a - -> Head abt a - -> Head abt a - -> Head abt a - -- TODO: something more efficient\/direct if we can... - evalOp And = \v1 v2 -> reflect (reify v1 && reify v2) - evalOp Or = \v1 v2 -> reflect (reify v1 || reify v2) - evalOp Xor = \v1 v2 -> reflect (reify v1 /= reify v2) - evalOp Iff = \v1 v2 -> reflect (reify v1 == reify v2) - evalOp (Min _) = error "TODO: evalOp{Min}" - evalOp (Max _) = error "TODO: evalOp{Max}" - {- - evalOp (Min _) = \v1 v2 -> reflect (reify v1 `min` reify v2) - evalOp (Max _) = \v1 v2 -> reflect (reify v1 `max` reify v2) - evalOp (Sum _) = \v1 v2 -> reflect (reify v1 + reify v2) - evalOp (Prod _) = \v1 v2 -> reflect (reify v1 * reify v2) - -} - -- HACK: this is just to have something to test. We really should reduce\/remove all this boilerplate... - evalOp (Sum theSemi) = - \(WLiteral v1) (WLiteral v2) -> WLiteral $ evalSum theSemi v1 v2 - evalOp (Prod theSemi) = - \(WLiteral v1) (WLiteral v2) -> WLiteral $ evalProd theSemi v1 v2 - - -- TODO: even if only one of the arguments is a literal, if that literal is zero\/one, then we can still partially evaluate it. (As is done in the old finally-tagless code) - evalSum, evalProd :: HSemiring a -> Literal a -> Literal a -> Literal a - evalSum HSemiring_Nat = \(LNat n1) (LNat n2) -> LNat (n1 + n2) - evalSum HSemiring_Int = \(LInt i1) (LInt i2) -> LInt (i1 + i2) - evalSum HSemiring_Prob = \(LProb p1) (LProb p2) -> LProb (p1 + p2) - evalSum HSemiring_Real = \(LReal r1) (LReal r2) -> LReal (r1 + r2) - evalProd HSemiring_Nat = \(LNat n1) (LNat n2) -> LNat (n1 * n2) - evalProd HSemiring_Int = \(LInt i1) (LInt i2) -> LInt (i1 * i2) - evalProd HSemiring_Prob = \(LProb p1) (LProb p2) -> LProb (p1 * p2) - evalProd HSemiring_Real = \(LReal r1) (LReal r2) -> LReal (r1 * r2) - -isIndex :: (ABT Term abt) => Variable 'HNat -> Dis abt Bool -isIndex v = do inds <- getIndices - return $ v `elem` map indVar inds - --- | For {evaluate, constrainValue v0} ArrayOp_ (Index _) :$ e1 :* e2 :* End -indexArrayOp :: forall abt typs args a r - . ( ABT Term abt - , typs ~ UnLCs args, args ~ LCs typs ) - => ArrayOp typs a - -> SArgs abt args - -> TermEvaluator abt (Dis abt) - -> (abt '[] a -> Dis abt r) -- evaluate or (constrainValue v0) - -> (Head abt ('HArray a) -- e1 is in whnf, and - -> Variable 'HNat -- e2 is a free under current indices - -> Dis abt r) - -> (Term abt ('HArray a) -- e1 is neutral syntax - -> Dis abt r) - -> (abt '[] ('HArray a) -- e1 is a free variable - -> Dis abt r) - -> (abt '[] ('HArray a) -- e1 is a multiloc, and - -> Variable 'HNat -- e2 is a free under current indices - -> Dis abt r) - -> Dis abt r -indexArrayOp o@(Index _) (e1 :* e2 :* End) evaluate_ kInd kArr kSyn kFree kMultiLoc = do - w1 <- evaluate_ e1 - case w1 of - Head_ arr@(WArray _ b) -> caseBind b $ \x body -> - evalIndex (kInd . flip (rename x) body) (kArr arr) - Head_ (WEmpty _) -> error "TODO: indexArrayOp o (Empty_ :* _ :* End)" - Head_ _ -> error "indexArrayOp: unknown whnf of array type" - Neutral e1' -> flip (caseVarSyn e1') kSyn $ \x -> - do locs <- getLocs - case (lookupAssoc x locs) of - Nothing -> kFree e1' - Just (Loc _ _) -> error "indexArrayOp: impossible, we have a Neutral term" - Just (MultiLoc l js) -> - evalIndex ((kInd . var =<<) . mkLoc Text.empty l . flip extendLocInds js) - (kMultiLoc e1') - where - evalIndex :: (ABT Term abt) - => (Variable 'HNat -> Dis abt r) - -> (Variable 'HNat -> Dis abt r) - -> Dis abt r - evalIndex thenCase elseCase = do - w2 <- evaluate_ e2 - caseWhnf w2 (const bot) $ \term -> -- bot if index is in whnf (eg. a literal num) - flip (caseVarSyn term) (const bot) $ \v -> -- bot if index is neutral syntax - do isInd <- isIndex v - if isInd then thenCase v else elseCase v -indexArrayOp _ _ _ _ _ _ _ _ = error "indexArrayOp called on incorrect ArrayOp" - -evaluateArrayOp - :: ( ABT Term abt - , typs ~ UnLCs args, args ~ LCs typs) - => TermEvaluator abt (Dis abt) - -> ArrayOp typs a - -> SArgs abt args - -> Dis abt (Whnf abt a) -evaluateArrayOp evaluate_ = go - where - go o@(Index _) = \args@(_ :* e2 :* End) -> - let returnIndex = return . Neutral . syn - in indexArrayOp o args - evaluate_ - evaluate_ - (\arr v -> returnIndex (ArrayOp_ o :$ fromHead arr :* var v :* End)) - (\s -> returnIndex (ArrayOp_ o :$ syn s :* e2 :* End)) - (\e1' -> returnIndex (ArrayOp_ o :$ e1' :* e2 :* End)) - (\e1' v -> returnIndex (ArrayOp_ o :$ e1' :* var v :* End)) - - go o@(Size _) = \(e1 :* End) -> do - w1 <- evaluate_ e1 - case w1 of - Neutral e1' -> return . Neutral $ syn (ArrayOp_ o :$ e1' :* End) - Head_ v1 -> - case head2array v1 of - WAEmpty -> return . Head_ $ WLiteral (LNat 0) - WAArray e3 _ -> evaluate_ e3 - - go (Reduce _) = \(e1 :* e2 :* e3 :* End) -> - error "TODO: evaluateArrayOp{Reduce}" - - -data ArrayHead :: ([Hakaru] -> Hakaru -> *) -> Hakaru -> * where - WAEmpty :: ArrayHead abt a - WAArray - :: !(abt '[] 'HNat) - -> !(abt '[ 'HNat] a) - -> ArrayHead abt a - -head2array :: Head abt ('HArray a) -> ArrayHead abt a -head2array (WEmpty _) = WAEmpty -head2array (WArray e1 e2) = WAArray e1 e2 - -impl, diff, nand, nor :: Bool -> Bool -> Bool -impl x y = not x || y -diff x y = x && not y -nand x y = not (x && y) -nor x y = not (x || y) - -evaluatePrimOp - :: forall abt p typs args a - . ( ABT Term abt - , typs ~ UnLCs args, args ~ LCs typs) - => TermEvaluator abt (Dis abt) - -> PrimOp typs a - -> SArgs abt args - -> Dis abt (Whnf abt a) -evaluatePrimOp evaluate_ = go - where - -- HACK: we don't have any way of saying these functions haven't reduced even though it's not actually a neutral term. - neu1 :: forall b c - . (abt '[] b -> abt '[] c) - -> abt '[] b - -> Dis abt (Whnf abt c) - neu1 f e = (Neutral . f . fromWhnf) <$> evaluate_ e - - neu2 :: forall b c d - . (abt '[] b -> abt '[] c -> abt '[] d) - -> abt '[] b - -> abt '[] c - -> Dis abt (Whnf abt d) - neu2 f e1 e2 = do e1' <- fromWhnf <$> evaluate_ e1 - e2' <- fromWhnf <$> evaluate_ e2 - return . Neutral $ f e1' e2' - - rr1 :: forall b b' c c' - . (Interp b b', Interp c c') - => (b' -> c') - -> (abt '[] b -> abt '[] c) - -> abt '[] b - -> Dis abt (Whnf abt c) - rr1 f' f e = do - w <- evaluate_ e - return $ - case w of - Neutral e' -> Neutral $ f e' - Head_ v -> Head_ . reflect $ f' (reify v) - - rr2 :: forall b b' c c' d d' - . (Interp b b', Interp c c', Interp d d') - => (b' -> c' -> d') - -> (abt '[] b -> abt '[] c -> abt '[] d) - -> abt '[] b - -> abt '[] c - -> Dis abt (Whnf abt d) - rr2 f' f e1 e2 = do - w1 <- evaluate_ e1 - w2 <- evaluate_ e2 - return $ - case w1 of - Neutral e1' -> Neutral $ f e1' (fromWhnf w2) - Head_ v1 -> - case w2 of - Neutral e2' -> Neutral $ f (fromWhnf w1) e2' - Head_ v2 -> Head_ . reflect $ f' (reify v1) (reify v2) - - primOp2_ - :: forall b c d - . PrimOp '[ b, c ] d -> abt '[] b -> abt '[] c -> abt '[] d - primOp2_ o e1 e2 = syn (PrimOp_ o :$ e1 :* e2 :* End) - - -- TODO: something more efficient\/direct if we can... - go Not (e1 :* End) = rr1 not P.not e1 - go Impl (e1 :* e2 :* End) = rr2 impl (primOp2_ Impl) e1 e2 - go Diff (e1 :* e2 :* End) = rr2 diff (primOp2_ Diff) e1 e2 - go Nand (e1 :* e2 :* End) = rr2 nand P.nand e1 e2 - go Nor (e1 :* e2 :* End) = rr2 nor P.nor e1 e2 - - -- HACK: we don't have a way of saying that 'Pi' (or 'Infinity',...) is in fact a head; so we're forced to call it neutral which is a lie. We should add constructor(s) to 'Head' to cover these magic constants; probably grouped together under a single constructor called something like @Constant@. Maybe should group them like that in the AST as well? - go Pi End = return $ Neutral P.pi - - -- We treat trig functions as strict, thus forcing their - -- arguments; however, to avoid fuzz issues we don't actually - -- evaluate the trig functions. - -- - -- HACK: we might should have some other way to make these - -- 'Whnf' rather than calling them neutral terms; since they - -- aren't, in fact, neutral! - go Sin (e1 :* End) = neu1 P.sin e1 - go Cos (e1 :* End) = neu1 P.cos e1 - go Tan (e1 :* End) = neu1 P.tan e1 - go Asin (e1 :* End) = neu1 P.asin e1 - go Acos (e1 :* End) = neu1 P.acos e1 - go Atan (e1 :* End) = neu1 P.atan e1 - go Sinh (e1 :* End) = neu1 P.sinh e1 - go Cosh (e1 :* End) = neu1 P.cosh e1 - go Tanh (e1 :* End) = neu1 P.tanh e1 - go Asinh (e1 :* End) = neu1 P.asinh e1 - go Acosh (e1 :* End) = neu1 P.acosh e1 - go Atanh (e1 :* End) = neu1 P.atanh e1 - - -- TODO: deal with how we have better types for these three ops than Haskell does... - -- go RealPow (e1 :* e2 :* End) = rr2 (**) (P.**) e1 e2 - go RealPow (e1 :* e2 :* End) = neu2 (P.**) e1 e2 - - -- HACK: these aren't actually neutral! - -- BUG: we should try to cancel out @(exp . log)@ and @(log . exp)@ - go Exp (e1 :* End) = neu1 P.exp e1 - go Log (e1 :* End) = neu1 P.log e1 - - -- HACK: these aren't actually neutral! - go (Infinity h) End = - case h of - HIntegrable_Nat -> return . Neutral $ P.primOp0_ (Infinity h) - HIntegrable_Prob -> return $ Neutral P.infinity - - go GammaFunc (e1 :* End) = neu1 P.gammaFunc e1 - go BetaFunc (e1 :* e2 :* End) = neu2 P.betaFunc e1 e2 - - go (Equal theEq) (e1 :* e2 :* End) = rrEqual theEq e1 e2 - go (Less theOrd) (e1 :* e2 :* End) = rrLess theOrd e1 e2 - go (NatPow theSemi) (e1 :* e2 :* End) = - case theSemi of - HSemiring_Nat -> rr2 (\v1 v2 -> v1 ^ fromNatural v2) (P.^) e1 e2 - HSemiring_Int -> rr2 (\v1 v2 -> v1 ^ fromNatural v2) (P.^) e1 e2 - HSemiring_Prob -> rr2 (\v1 v2 -> v1 ^ fromNatural v2) (P.^) e1 e2 - HSemiring_Real -> rr2 (\v1 v2 -> v1 ^ fromNatural v2) (P.^) e1 e2 - go (Negate theRing) (e1 :* End) = - case theRing of - HRing_Int -> rr1 negate P.negate e1 - HRing_Real -> rr1 negate P.negate e1 - go (Abs theRing) (e1 :* End) = - case theRing of - HRing_Int -> rr1 (unsafeNatural . abs) P.abs_ e1 - HRing_Real -> rr1 (unsafeNonNegativeRational . abs) P.abs_ e1 - go (Signum theRing) (e1 :* End) = - case theRing of - HRing_Int -> rr1 signum P.signum e1 - HRing_Real -> rr1 signum P.signum e1 - go (Recip theFractional) (e1 :* End) = - case theFractional of - HFractional_Prob -> rr1 recip P.recip e1 - HFractional_Real -> rr1 recip P.recip e1 - go (NatRoot theRadical) (e1 :* e2 :* End) = - case theRadical of - HRadical_Prob -> neu2 (flip P.thRootOf) e1 e2 - {- - go (NatRoot theRadical) (e1 :* e2 :* End) = - case theRadical of - HRadical_Prob -> rr2 natRoot (flip P.thRootOf) e1 e2 - go (Erf theContinuous) (e1 :* End) = - case theContinuous of - HContinuous_Prob -> rr1 erf P.erf e1 - HContinuous_Real -> rr1 erf P.erf e1 - -} - go op _ = error $ "TODO: evaluatePrimOp{" ++ show op ++ "}" - - - rrEqual - :: forall b. HEq b -> abt '[] b -> abt '[] b -> Dis abt (Whnf abt HBool) - rrEqual theEq = - case theEq of - HEq_Nat -> rr2 (==) (P.==) - HEq_Int -> rr2 (==) (P.==) - HEq_Prob -> rr2 (==) (P.==) - HEq_Real -> rr2 (==) (P.==) - HEq_Array aEq -> error "TODO: rrEqual{HEq_Array}" - HEq_Bool -> rr2 (==) (P.==) - HEq_Unit -> rr2 (==) (P.==) - HEq_Pair aEq bEq -> - \e1 e2 -> do - w1 <- evaluate_ e1 - w2 <- evaluate_ e2 - case w1 of - Neutral e1' -> - return . Neutral - $ P.primOp2_ (Equal theEq) e1' (fromWhnf w2) - Head_ v1 -> - case w2 of - Neutral e2' -> - return . Neutral - $ P.primOp2_ (Equal theEq) (fromHead v1) e2' - Head_ v2 -> do - let (v1a, v1b) = reifyPair v1 - let (v2a, v2b) = reifyPair v2 - wa <- rrEqual aEq v1a v2a - wb <- rrEqual bEq v1b v2b - return $ - case wa of - Neutral ea -> - case wb of - Neutral eb -> Neutral (ea P.&& eb) - Head_ vb - | reify vb -> wa - | otherwise -> Head_ $ WDatum dFalse - Head_ va - | reify va -> wb - | otherwise -> Head_ $ WDatum dFalse - - HEq_Either aEq bEq -> error "TODO: rrEqual{HEq_Either}" - - rrLess - :: forall b. HOrd b -> abt '[] b -> abt '[] b -> Dis abt (Whnf abt HBool) - rrLess theOrd = - case theOrd of - HOrd_Nat -> rr2 (<) (P.<) - HOrd_Int -> rr2 (<) (P.<) - HOrd_Prob -> rr2 (<) (P.<) - HOrd_Real -> rr2 (<) (P.<) - HOrd_Array aOrd -> error "TODO: rrLess{HOrd_Array}" - HOrd_Bool -> rr2 (<) (P.<) - HOrd_Unit -> rr2 (<) (P.<) - HOrd_Pair aOrd bOrd -> - \e1 e2 -> do - w1 <- evaluate_ e1 - w2 <- evaluate_ e2 - case w1 of - Neutral e1' -> - return . Neutral - $ P.primOp2_ (Less theOrd) e1' (fromWhnf w2) - Head_ v1 -> - case w2 of - Neutral e2' -> - return . Neutral - $ P.primOp2_ (Less theOrd) (fromHead v1) e2' - Head_ v2 -> do - let (v1a, v1b) = reifyPair v1 - let (v2a, v2b) = reifyPair v2 - error "TODO: rrLess{HOrd_Pair}" - -- BUG: The obvious recursion won't work because we need to know when the first components are equal before recursing (to implement lexicographic ordering). We really need a ternary comparison operator like 'compare'. - HOrd_Either aOrd bOrd -> error "TODO: rrLess{HOrd_Either}" - -update - :: forall abt - . (ABT Term abt) - => MeasureEvaluator abt (Dis abt) - -> TermEvaluator abt (Dis abt) - -> VariableEvaluator abt (Dis abt) -update perform evaluate_ x = - do locs <- getLocs - -- If we get 'Nothing', then it turns out @x@ is a free variable - maybe (return $ Neutral (var x)) lookForLoc (lookupAssoc x locs) - where lookForLoc (Loc l jxs) = - (maybe (freeLocError l) return =<<) . select l $ \s -> - case s of - SBind l' e ixs -> do - Refl <- varEq l l' - Just $ do - w <- withIndices ixs $ perform (caseLazy e fromWhnf id) - unsafePush (SLet l (Whnf_ w) ixs) -#ifdef __TRACE_DISINTEGRATE__ - trace ("-- updated " - ++ show (ppStatement 11 s) - ++ " to " - ++ show (ppStatement 11 (SLet l (Whnf_ w) ixs)) - ) $ return () -#endif - let as = toAssocs $ zipWith Assoc (map indVar ixs) jxs - w' = renames as (fromWhnf w) - inds <- getIndices - withIndices inds $ return (fromMaybe (Neutral w') (toWhnf w')) - SLet l' e ixs -> do - Refl <- varEq l l' - Just $ do - w <- withIndices ixs $ caseLazy e return evaluate_ - unsafePush (SLet l (Whnf_ w) ixs) - let as = toAssocs $ zipWith Assoc (map indVar ixs) jxs - w' = renames as (fromWhnf w) - inds <- getIndices - withIndices inds $ return (fromMaybe (Neutral w') (toWhnf w')) - -- This does not bind any variables, so it definitely can't match. - SWeight _ _ -> Nothing - -- This does bind variables, - -- but there's no expression we can return for it - -- because the variables are untouchable\/abstract. - SGuard ls pat scrutinee i -> Just . return . Neutral $ var x - - -- Case for MultiLocs - lookForLoc (MultiLoc l jxs) = return (Neutral $ var x) - ----------------------------------------------------------- End of copied code -- - - --- | The forward disintegrator's function for evaluating case --- expressions. First we try calling 'defaultCaseEvaluator' which --- will evaluate the scrutinee and select the matching branch (if --- any). But that doesn't work out in general, since the scrutinee --- may contain heap-bound variables. So our fallback definition --- will push a 'SGuard' onto the heap and then continue evaluating --- each branch (thereby duplicating the continuation, calling it --- once on each branch). -evaluateCase - :: forall abt - . (ABT Term abt) - => TermEvaluator abt (Dis abt) - -> CaseEvaluator abt (Dis abt) -{-# INLINE evaluateCase #-} -evaluateCase evaluate_ = evaluateCase_ - where - evaluateCase_ :: CaseEvaluator abt (Dis abt) - evaluateCase_ e bs = - defaultCaseEvaluator evaluate_ e bs - <|> evaluateBranches e bs - - evaluateBranches :: CaseEvaluator abt (Dis abt) - evaluateBranches e = choose . map evaluateBranch - where - evaluateBranch (Branch pat body) = - let (vars,body') = caseBinds body - in getIndices >>= \i -> - push (SGuard vars pat (Thunk e) i) body' evaluate_ - +evaluate_ = evaluate perform evaluateDatum :: (ABT Term abt) => DatumEvaluator (abt '[]) (Dis abt) evaluateDatum e = viewWhnfDatum <$> evaluate_ e --- | Simulate performing 'HMeasure' actions by simply emiting code +-- | Simulate performing 'HMeasure' actions by simply emitting code -- for those actions, returning the bound variable. -- -- This is the function called @(|>>)@ in the disintegration paper. @@ -920,11 +290,11 @@ perform = \e0 -> #ifdef __TRACE_DISINTEGRATE__ getStatements >>= \ss -> - getLocs >>= \locs -> + getExtras >>= \extras -> getIndices >>= \inds -> trace ("\n-- perform --\n" ++ "at " ++ show (ppInds inds) ++ "\n" - ++ show (prettyLocs locs) ++ "\n" + ++ show (prettyExtras extras) ++ "\n" ++ show (pretty_Statements_withTerm ss e0) ++ "\n") $ #endif @@ -936,35 +306,29 @@ performTerm (MBind :$ e1 :* e2 :* End) = caseBind e2 $ \x e2' -> do inds <- getIndices - push (SBind x (Thunk e1) inds) e2' perform + push (SBind x (Thunk e1) inds) e2' >>= perform performTerm (Plate :$ e1 :* e2 :* End) = do x1 <- pushPlate e1 e2 - return (Neutral (var x1)) + return $ fromJust (toWhnf x1) performTerm (Superpose_ pes) = do inds <- getIndices if not (null inds) && L.length pes > 1 then bot else emitFork_ (P.superpose . fmap ((,) P.one)) - (fmap (\(p,e) -> push (SWeight (Thunk p) inds) e perform) + (fmap (\(p,e) -> push (SWeight (Thunk p) inds) e >>= perform) pes) -- Avoid falling through to the @performWhnf <=< evaluate_@ case performTerm (Let_ :$ e1 :* e2 :* End) = caseBind e2 $ \x e2' -> do inds <- getIndices - push (SLet x (Thunk e1) inds) e2' perform + push (SLet x (Thunk e1) inds) e2' >>= perform -- TODO: we could optimize this by calling some @evaluateTerm@ -- directly, rather than calling 'syn' to rebuild @e0@ from -- @t0@ and then calling 'evaluate_' (which will just use -- 'caseVarSyn' to get the @t0@ back out from the @e0@). - -- - -- BUG: when @t0@ is a 'Case_', this doesn't work right. This - -- is the source of the hygiene bug in 'testPerform1b'. Alas, - -- we cannot use 'emitCaseWith' here since that would require - -- the scrutinee to be emissible; but we'd want something pretty - -- similar... performTerm t0 = do w <- evaluate_ (syn t0) #ifdef __TRACE_DISINTEGRATE__ @@ -974,15 +338,13 @@ performVar :: forall a. Variable ('HMeasure a) -> Dis abt (Whnf abt a) - performVar = performWhnf <=< update perform evaluate_ + performVar = performWhnf <=< evaluateVar perform evaluate_ - -- BUG: it's not clear this is actually doing the right thing for its call-sites. In particular, we should handle 'Case_' specially, to deal with the hygiene bug in 'testPerform1b'... - -- - -- BUG: found the 'testPerform1b' hygiene bug! We can't simply call 'emitMBind' on @e@, because @e@ may not be emissible! performWhnf :: forall a. Whnf abt ('HMeasure a) -> Dis abt (Whnf abt a) performWhnf (Head_ v) = perform $ fromHead v - performWhnf (Neutral e) = (Neutral . var) <$> emitMBind e + performWhnf (Neutral e) = (Neutral . var) <$> + (emitMBind . fromWhnf =<< atomize e) -- TODO: right now we do the simplest thing. However, for better @@ -1005,10 +367,12 @@ -> Dis abt (Whnf abt a) nice o es = do es' <- traverse21 atomizeCore es - x <- emitMBind $ syn (MeasureOp_ o :$ es') - return (Neutral $ var x) + x <- emitMBind2 $ syn (MeasureOp_ o :$ es') + -- return (Neutral $ var x) + return (Neutral x) - -- Try to be as complete as possible (i.e., 'bot' as little as possible), no matter how ugly the output code gets. + -- Try to be as complete as possible (i.e., 'bot' as little as + -- possible), no matter how ugly the output code gets. complete :: MeasureOp typs a -> SArgs abt args @@ -1023,18 +387,6 @@ pushWeight (P.densityUniform lo hi x) return (Neutral x) complete _ = \_ -> bot - --- Calls unsafePush -pushWeight :: (ABT Term abt) => abt '[] 'HProb -> Dis abt () -pushWeight w = do - inds <- getIndices - unsafePush $ SWeight (Thunk w) inds - --- Calls unsafePush -pushGuard :: (ABT Term abt) => abt '[] HBool -> Dis abt () -pushGuard b = do - inds <- getIndices - unsafePush $ SGuard Nil1 pTrue (Thunk b) inds -- | The goal of this function is to ensure the correctness criterion @@ -1060,7 +412,11 @@ #ifdef __TRACE_DISINTEGRATE__ trace ("\n-- atomize --\n" ++ show (pretty e)) $ #endif - traverse21 atomizeCore =<< evaluate_ e + do whnf <- evaluate_ e + case whnf of + Head_ v -> Head_ <$> traverse21 atomizeCore v + Neutral e -> Neutral . unviewABT <$> + traverse12 (traverse21 atomizeCore) (viewABT e) -- | A variant of 'atomize' which is polymorphic in the locally @@ -1068,20 +424,25 @@ -- We factored this out because we often want this more polymorphic -- variant when using our indexed @TraversableMN@ classes. atomizeCore :: (ABT Term abt) => abt xs a -> Dis abt (abt xs a) -atomizeCore e = do +atomizeCore e = -- HACK: this check for 'disjointVarSet' is sufficient to catch -- the particular infinite loops we were encountering, but it -- will not catch all of them. If the call to 'evaluate_' in -- 'atomize' returns a neutral term which contains heap-bound -- variables, then we'll still loop forever since we don't -- traverse\/fmap over the top-level term constructor of neutral - -- terms. - xs <- getHeapVars - if disjointVarSet xs (freeVars e) + -- terms. + do xs <- getHeapVars + vs <- extFreeVars e + if disjointVarSet xs vs then return e else let (ys, e') = caseBinds e - in (binds_ ys . fromWhnf) <$> atomize e' + in +#ifdef __TRACE_DISINTEGRATE__ + trace ("\n-- atomizeCore --\n" ++ show (pretty e')) $ +#endif + (binds_ ys . fromWhnf) <$> atomize e' where -- TODO: does @IM.null . IM.intersection@ fuse correctly? disjointVarSet xs ys = @@ -1095,7 +456,6 @@ Dis $ \_ c h -> c (foldMap statementVars (statements h)) h ---------------------------------------------------------------- ----------------------------------------------------------------- -- | Given an emissible term @v0@ (the first argument) and another -- term @e0@ (the second argument), compute the constraints such -- that @e0@ must evaluate to @v0@. This is the function called @@ -1113,35 +473,28 @@ constrainValue v0 e0 = #ifdef __TRACE_DISINTEGRATE__ getStatements >>= \ss -> - getLocs >>= \locs -> + getExtras >>= \extras -> getIndices >>= \inds -> trace ("\n-- constrainValue: " ++ show (pretty v0) ++ "\n" ++ show (pretty_Statements_withTerm ss e0) ++ "\n" ++ "at " ++ show (ppInds inds) ++ "\n" - ++ show (prettyLocs locs) ++ "\n" + ++ show (prettyExtras extras) ++ "\n" ) $ #endif caseVarSyn e0 (constrainVariable v0) $ \t -> case t of -- There's a bunch of stuff we don't even bother trying to handle - Empty_ _ -> error "TODO: disintegrate arrays" + Empty_ _ -> error "TODO: disintegrate empty arrays" Array_ n e -> caseBind e $ \x body -> do j <- freshInd n let x' = indVar j - body' = rename x x' body - inds <- getIndices + body' <- extSubst x (var x') body + inds <- getIndices withIndices (extendIndices j inds) $ constrainValue (v0 P.! (var x')) body' -- TODO use meta-index - ArrayOp_ o@(Index _) :$ args -> indexArrayOp o args - evaluate_ - (constrainValue v0) - (const $ const bot) - (const bot) - (const bot) - (const $ const bot) - - ArrayOp_ _ :$ _ -> error "TODO: disintegrate arrays" + ArrayLiteral_ _ -> error "TODO: disintegrate literal arrays" + ArrayOp_ o :$ args -> constrainValueArrayOp v0 o args Lam_ :$ _ :* End -> error "TODO: disintegrate lambdas" App_ :$ _ :* _ :* End -> error "TODO: disintegrate lambdas" Integrate :$ _ :* _ :* _ :* End -> @@ -1169,7 +522,7 @@ Reject_ _ -> bot -- giving up. Let_ :$ e1 :* e2 :* End -> caseBind e2 $ \x e2' -> - push (SLet x (Thunk e1) []) e2' (constrainValue v0) + push (SLet x (Thunk e1) []) e2' >>= constrainValue v0 CoerceTo_ c :$ e1 :* End -> -- TODO: we need to insert some kind of guard that says @@ -1211,7 +564,7 @@ Just GotStuck -> constrainBranches v0 e bs Just (Matched rho body) -> - pushes (toStatements rho) body (constrainValue v0) + pushes (toVarStatements rho) body >>= constrainValue v0 <|> constrainBranches v0 e bs _ :$ _ -> error "constrainValue: the impossible happened" @@ -1238,7 +591,8 @@ where constrainBranch (Branch pat body) = let (vars,body') = caseBinds body - in push (SGuard vars pat (Thunk e) []) body' (constrainValue v0) + in push (SGuard vars pat (Thunk e) []) body' + >>= constrainValue v0 constrainDatum @@ -1313,53 +667,39 @@ constrainVariable :: (ABT Term abt) => abt '[] a -> Variable a -> Dis abt () constrainVariable v0 x = - do locs <- getLocs + do extras <- getExtras -- If we get 'Nothing', then it turns out @x@ is a free variable. -- If @x@ is a free variable, then it's a neutral term; and we -- return 'bot' for neutral terms - maybe bot lookForLoc (lookupAssoc x locs) + maybe bot lookForLoc (lookupAssoc x extras) where lookForLoc (Loc l jxs) = - let -- Assumption: js has no duplicates - permutes is js = length is == length js && - Set.fromList is == Set.fromList (map indVar js) - -- If we get 'Nothing', then it turns out @l@ is a free location. - -- This is an error because of the invariant: - -- if there exists an 'Assoc x ({Multi}Loc l _)' inside @locs@ + -- If we get 'Nothing', then it turns out @l@ is a free + -- location. This is an error because of the + -- invariant: + -- if there exists an 'Assoc x (Loc l _)' inside @extras@ -- then there must be a statement on the 'ListContext' that binds @l@ - in (maybe (freeLocError l) return =<<) . select l $ \s -> + (maybe (freeLocError l) return =<<) . select l $ \s -> case s of SBind l' e ixs -> do - Refl <- varEq l l' + Refl <- locEq l l' guard (length ixs == length jxs) -- will error otherwise Just $ do inds <- getIndices guard (jxs `permutes` inds) -- will bot otherwise - e' <- apply (zip ixs inds) (fromLazy e) + e' <- apply ixs inds (fromLazy e) constrainOutcome v0 e' unsafePush (SLet l (Whnf_ (Neutral v0)) inds) SLet l' e ixs -> do - Refl <- varEq l l' + Refl <- locEq l l' guard (length ixs == length jxs) -- will error otherwise Just $ do inds <- getIndices guard (jxs `permutes` inds) -- will bot otherwise - e' <- apply (zip ixs inds) (fromLazy e) + e' <- apply ixs inds (fromLazy e) constrainValue v0 e' unsafePush (SLet l (Whnf_ (Neutral v0)) inds) SWeight _ _ -> Nothing SGuard ls' pat scrutinee i -> error "TODO: constrainVariable{SGuard}" - -- Case for MultiLoc - lookForLoc (MultiLoc l jxs) = do -#ifdef __TRACE_DISINTEGRATE__ - traceM $ "looking for MultiLoc: " ++ show (prettyLoc (MultiLoc l jxs)) -#endif - n <- sizeInnermostInd l - j <- freshInd n - x' <- mkLoc Text.empty l (extendLocInds (indVar j) jxs) - inds <- getIndices - withIndices (extendIndices j inds) $ - constrainValue (v0 P.! (var $ indVar j)) (var x') - -- TODO use meta-index ---------------------------------------------------------------- -- | N.B., as with 'constrainValue', we assume that the first @@ -1368,6 +708,82 @@ -- -- TODO: capture the emissibility requirement on the first argument -- in the types. +constrainValueArrayOp + :: forall abt typs args a + . (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) + => abt '[] a + -> ArrayOp typs a + -> SArgs abt args + -> Dis abt () +constrainValueArrayOp v0 = go + where + go :: ArrayOp typs a -> SArgs abt args -> Dis abt () + go (Index _) (e1 :* e2 :* End) = do + w1 <- evaluate_ e1 + case w1 of + Neutral e1' -> bot + Head_ (WArray _ b) -> caseBind b $ \x body -> + extSubst x e2 body >>= constrainValue v0 + Head_ (WEmpty _) -> bot -- TODO: check this + Head_ a@(WArrayLiteral _) -> constrainValueIdxArrLit v0 e2 a + _ -> error "constrainValue {ArrayOp Index}: uknown whnf of array type" + go (Size _) _ = error "TODO: disintegrate {ArrayOp Size}" + go (Reduce _) _ = error "TODO: disintegrate {ArrayOp Reduce}" + go _ _ = error "constrainValueArrayOp: unknown arrayOp" + + +-- | Special case for [true,false] ! i, and [false, true] ! i +-- This helps us disintegrate bern +constrainValueIdxArrLit + :: forall abt a + . (ABT Term abt) + => abt '[] a + -> abt '[] 'HNat + -> Head abt ('HArray a) + -> Dis abt () +constrainValueIdxArrLit v0 e2 = go + where + go :: Head abt ('HArray a) -> Dis abt () + go (WArrayLiteral [a1,a2]) = + case (jmEq1 (typeOf v0) sBool) of + Just Refl -> + let constrainInd = flip constrainValue e2 + in case (isLitBool a1, isLitBool a2) of + (Just b1, Just b2) + | isLitTrue b1 && isLitFalse b2 -> + constrainInd $ P.if_ v0 (P.nat_ 0) (P.nat_ 1) + | isLitTrue b2 && isLitFalse b1 -> + constrainInd $ P.if_ v0 (P.nat_ 1) (P.nat_ 0) + | otherwise -> error "constrainValue: cannot invert (Index [b,b] i)" + _ -> error "TODO: constrainValue (Index [b1,b2] i)" + Nothing -> error "TODO: constrainValue (Index [a1,a2] i)" + go (WArrayLiteral _) = bot + go _ = error "constrainValueIdxArrLit: unknown ArrayLiteral form" + +-- | Helpers for disintegrating bern +isLitBool :: (ABT Term abt) => abt '[] a -> Maybe (Datum (abt '[]) HBool) +isLitBool e = caseVarSyn e (const Nothing) $ \b -> + case b of + Datum_ d@(Datum _ typ _) -> case (jmEq1 typ sBool) of + Just Refl -> Just d + Nothing -> Nothing + _ -> Nothing + +isLitTrue :: (ABT Term abt) => Datum (abt '[]) HBool -> Bool +isLitTrue (Datum tdTrue sBool (Inl Done)) = True +isLitTrue _ = False + +isLitFalse :: (ABT Term abt) => Datum (abt '[]) HBool -> Bool +isLitFalse (Datum tdFalse sBool (Inr (Inl Done))) = True +isLitFalse _ = False + +---------------------------------------------------------------- +-- | N.B., as with 'constrainValue', we assume that the first +-- argument is emissible. So it is the caller's responsibility to +-- ensure this (by calling 'atomize' as appropriate). +-- +-- TODO: capture the emissibility requirement on the first argument +-- in the types. constrainValueMeasureOp :: forall abt typs args a . (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) @@ -1458,8 +874,11 @@ toProb_abs HSemiring_Real = P.abs_ --- TODO: is there any way to optimise the zippering over the Seq, a la 'S.inits' or 'S.tails'? --- TODO: really we want a dynamic programming approach to avoid unnecessary repetition of calling @evaluateNaryOp evaluate_@ on the two subsequences... +-- TODO: is there any way to optimise the zippering over the Seq, a la +-- 'S.inits' or 'S.tails'? +-- TODO: really we want a dynamic programming +-- approach to avoid unnecessary repetition of calling @evaluateNaryOp +-- evaluate_@ on the two subsequences... lubSeq :: (Alternative m) => (Seq a -> a -> Seq a -> m b) -> Seq a -> m b lubSeq f = go S.empty where @@ -1572,9 +991,11 @@ go Acosh = \(e1 :* End) -> error_TODO "Acosh" go Atanh = \(e1 :* End) -> error_TODO "Atanh" go RealPow = \(e1 :* e2 :* End) -> - -- TODO: There's a discrepancy between @(**)@ and @pow_@ in the old code... + -- TODO: There's a discrepancy between @(**)@ and @pow_@ in + -- the old code... do - -- TODO: if @v1@ is 0 or 1 then bot. Maybe the @log v1@ in @w@ takes care of the 0 case? + -- TODO: if @v1@ is 0 or 1 then bot. Maybe the @log v1@ in + -- @w@ takes care of the 0 case? u <- emitLet' v0 -- either this from @(**)@: -- emitGuard $ P.zero P.< u @@ -1587,7 +1008,8 @@ constrainValue (P.log u P./ P.log e1) e2 -- end. <|> do - -- TODO: if @v2@ is 0 then bot. Maybe the weight @w@ takes care of this case? + -- TODO: if @v2@ is 0 then bot. Maybe the weight @w@ takes + -- care of this case? u <- emitLet' v0 let ex = v0 P.** P.recip e2 -- either this from @(**)@: @@ -1600,7 +1022,9 @@ constrainValue ex e1 go Exp = \(e1 :* End) -> do x0 <- emitLet' v0 - -- TODO: do we still want\/need the @emitGuard (0 < x0)@ which is now equivalent to @emitGuard (0 /= x0)@ thanks to the types? + -- TODO: do we still want\/need the @emitGuard (0 < x0)@ which + -- is now equivalent to @emitGuard (0 /= x0)@ thanks to the + -- types? emitWeight (P.recip x0) constrainValue (P.log x0) e1 @@ -1616,8 +1040,11 @@ go (Less theOrd) = \(e1 :* e2 :* End) -> error_TODO "Less" go (NatPow theSemi) = \(e1 :* e2 :* End) -> error_TODO "NatPow" go (Negate theRing) = \(e1 :* End) -> - -- TODO: figure out how to merge this implementation of @rr1 negate@ with the one in 'evaluatePrimOp' to DRY - -- TODO: just emitLet the @v0@ and pass the neutral term to the recursive call? + -- TODO: figure out how to merge this implementation of @rr1 + -- negate@ with the one in 'evaluatePrimOp' to DRY + -- TODO: just + -- emitLet the @v0@ and pass the neutral term to the recursive + -- call? let negate_v0 = syn (PrimOp_ (Negate theRing) :$ v0 :* End) -- case v0 of -- Neutral e -> @@ -1663,7 +1090,8 @@ emitWeight . P.recip . P.unsafeProbFraction_ theFractional - -- TODO: define a dictionary-passing variant of 'P.square' instead, to include the coercion in there explicitly... + -- TODO: define a dictionary-passing variant of 'P.square' + -- instead, to include the coercion in there explicitly... $ square (hSemiring_HFractional theFractional) x0 constrainValue (P.primOp1_ (Recip theFractional) x0) e1 @@ -1720,7 +1148,7 @@ -> Dis abt () constrainOutcome v0 e0 = #ifdef __TRACE_DISINTEGRATE__ - getLocs >>= \locs -> + getExtras >>= \extras -> getIndices >>= \inds -> trace ( let s = "-- constrainOutcome" @@ -1729,7 +1157,7 @@ ++ "\n" ++ replicate (length s) ' ' ++ ": " ++ show (pretty e0) ++ "\n" ++ "at " ++ show (ppInds inds) ++ "\n" - ++ show (prettyLocs locs) + ++ show (prettyExtras extras) ) $ #endif do w0 <- evaluate_ e0 @@ -1754,10 +1182,10 @@ go (WMBind e1 e2) = caseBind e2 $ \x e2' -> do i <- getIndices - push (SBind x (Thunk e1) i) e2' (constrainOutcome v0) + push (SBind x (Thunk e1) i) e2' >>= constrainOutcome v0 go (WPlate e1 e2) = do x' <- pushPlate e1 e2 - constrainValue v0 (var x') + constrainValue v0 x' go (WChain e1 e2 e3) = error "TODO: constrainOutcome{Chain}" go (WReject typ) = emit_ $ \m -> P.reject (typeOf m) @@ -1765,7 +1193,7 @@ i <- getIndices if not (null i) && L.length pes > 1 then bot else emitFork_ (P.superpose . fmap ((,) P.one)) - (fmap (\(p,e) -> push (SWeight (Thunk p) i) e (constrainOutcome v0)) + (fmap (\(p,e) -> push (SWeight (Thunk p) i) e >>= constrainOutcome v0) pes) -- TODO: should this really be different from 'constrainValueMeasureOp'? @@ -1804,7 +1232,12 @@ -- y <~ normal(x,1) -- return ((x+(y+y),x)::pair(real,real)) go Normal = \(mu :* sd :* End) -> do - -- N.B., if\/when extending this to higher dimensions, the real equation is @recip (sqrt (2*pi*sd^2) ^ n) * exp (negate (norm_n (v0 - mu) ^ 2) / (2*sd^2))@ for @Real^n@. + -- N.B., if\/when extending this to higher dimensions, the + -- real equation is + -- @recip (sqrt (2*pi*sd^2) ^ n) * + -- exp (negate (norm_n (v0 - mu) ^ 2) / + -- (2*sd^2))@ + -- for @Real^n@. pushWeight (P.densityNormal mu sd v0) go Poisson = \(e1 :* End) -> do diff --git a/haskell/Language/Hakaru/Evaluation/ConstantPropagation.hs b/haskell/Language/Hakaru/Evaluation/ConstantPropagation.hs --- a/haskell/Language/Hakaru/Evaluation/ConstantPropagation.hs +++ b/haskell/Language/Hakaru/Evaluation/ConstantPropagation.hs @@ -1,12 +1,13 @@ -{-# LANGUAGE CPP - , GADTs - , DataKinds - , Rank2Types - , ScopedTypeVariables - , MultiParamTypeClasses - , FlexibleContexts - , FlexibleInstances - #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_GHC -Wall -fwarn-tabs #-} ---------------------------------------------------------------- @@ -26,43 +27,91 @@ ) where #if __GLASGOW_HASKELL__ < 710 -import Data.Functor ((<$>)) -import Control.Applicative (Applicative(..)) +import Control.Applicative (Applicative (..)) +import Data.Functor ((<$>)) #endif -import Language.Hakaru.Syntax.IClasses (Traversable21(..)) -import Language.Hakaru.Syntax.ABT (View(..), ABT(..), cataABT) -import Language.Hakaru.Syntax.AST -import Language.Hakaru.Evaluation.EvalMonad (runPureEvaluate) +import Control.Monad.Reader +import Data.Monoid (All (..)) +import Language.Hakaru.Evaluation.EvalMonad (runPureEvaluate) +import Language.Hakaru.Syntax.ABT (ABT (..), View (..)) +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.IClasses (Foldable21 (..), + Traversable21 (..)) +import Language.Hakaru.Syntax.Variable +type Env = Assocs Literal + +-- The constant propagation monad. Simply threads through an environment mapping +-- variables to known constant values. +newtype PropM a = PropM { runPropM :: Reader Env a } + deriving (Functor, Applicative, Monad, MonadReader Env) + ---------------------------------------------------------------- ---------------------------------------------------------------- -- TODO: try evaluating certain things even if not all their immediate -- subterms are literals. For example: --- (1) substitute let-bindings of literals --- (2) evaluate beta-redexes where the argument is a literal --- (3) evaluate case-of-constructor if we can --- (4) handle identity elements for NaryOps --- N.B., some of these will require performing top-down work to --- avoid re-traversing subtrees. +-- (1) evaluate beta-redexes where the argument is a literal +-- (2) evaluate case-of-constructor if we can +-- (3) handle identity elements for NaryOps +-- (4) Recognize trivial cases for looping constructs: +-- summate a b (const 0) == 0 +-- summate a b id == b - a +-- summate a b (const x) == x * (b - a) -- -- | Perform basic constant propagation. constantPropagation - :: forall abt a. (ABT Term abt) => abt '[] a -> abt '[] a -constantPropagation = - cataABT var bind alg - where - getLiteral :: forall xs b. abt xs b -> Maybe (abt xs b) - getLiteral e = - case viewABT e of - Syn (Literal_ _) -> Just e - _ -> Nothing + :: forall abt a . (ABT Term abt) + => abt '[] a + -> abt '[] a +constantPropagation abt = runReader (runPropM $ constantProp' abt) emptyAssocs - alg :: forall b. Term abt b -> abt '[] b - alg t = - case traverse21 getLiteral t of - Nothing -> syn t - Just t' -> runPureEvaluate (syn t') +constantProp' + :: forall abt a xs . (ABT Term abt) + => abt xs a + -> PropM (abt xs a) +constantProp' = start + where ----------------------------------------------------------------- ------------------------------------------------------------ fin. + start :: forall b ys . abt ys b -> PropM (abt ys b) + start = loop . viewABT + + loop :: forall b ys . View (Term abt) ys b -> PropM (abt ys b) + loop (Var v) = maybe (var v) (syn . Literal_) . lookupAssoc v <$> ask + loop (Syn s) = constantPropTerm s + loop (Bind v b) = bind v <$> loop b + +isLiteral :: forall abt b ys . (ABT Term abt) => abt ys b -> Bool +isLiteral abt = case viewABT abt of + Syn (Literal_ _) -> True + _ -> False + +isFoldable :: forall abt b . (ABT Term abt) => Term abt b -> Bool +isFoldable = getAll . foldMap21 (All . isLiteral) + +getLiteral :: forall abt ys b. (ABT Term abt) => abt ys b -> Maybe (Literal b) +getLiteral e = + case viewABT e of + Syn (Literal_ l) -> Just l + _ -> Nothing + +tryEval :: forall abt b . (ABT Term abt) => Term abt b -> abt '[] b +tryEval term + | isFoldable term = runPureEvaluate (syn term) + | otherwise = syn term + +constantPropTerm + :: (ABT Term abt) + => Term abt a + -> PropM (abt '[] a) +constantPropTerm (Let_ :$ rhs :* body :* End) = + caseBind body $ \v body' -> do + rhs' <- constantProp' rhs + case getLiteral rhs' of + Just l -> local (insertAssoc (Assoc v l)) (constantProp' body') + Nothing -> do + body'' <- constantProp' body' + return $ syn (Let_ :$ rhs' :* bind v body'' :* End) + +constantPropTerm term = tryEval <$> traverse21 constantProp' term + diff --git a/haskell/Language/Hakaru/Evaluation/DisintegrationMonad.hs b/haskell/Language/Hakaru/Evaluation/DisintegrationMonad.hs --- a/haskell/Language/Hakaru/Evaluation/DisintegrationMonad.hs +++ b/haskell/Language/Hakaru/Evaluation/DisintegrationMonad.hs @@ -40,7 +40,7 @@ --, reject -- ** Emitting code , emit - , emitMBind + , emitMBind , emitMBind2 , emitLet , emitLet' , emitUnpair @@ -53,27 +53,31 @@ , emitFork_ , emitSuperpose , choose + , pushWeight + , pushGuard -- * Overrides for original in Evaluation.Types , pushPlate -- * For Arrays/Plate , getIndices , withIndices , extendIndices - , extendLocInds + , selectMore + , permutes , statementInds , sizeInnermostInd - -- * Locs - , Loc(..) - , getLocs - , putLocs - , insertLoc - , adjustLoc + -- * Extras + , Extra(..) + , getExtras + , putExtras + , insertExtra + , adjustExtra , mkLoc , freeLocError - , apply + , zipInds + , apply #ifdef __TRACE_DISINTEGRATE__ - , prettyLoc - , prettyLocs + , prettyExtra + , prettyExtras #endif ) where @@ -84,6 +88,7 @@ import Data.Functor ((<$>)) import Control.Applicative (Applicative(..)) #endif +import qualified Data.Set as Set (fromList) import Data.Maybe import qualified Data.Foldable as F import qualified Data.Traversable as T @@ -114,13 +119,13 @@ import Language.Hakaru.Pretty.Haskell (ppVariable, pretty) #endif -getStatements :: Dis abt [Statement abt 'Impure] +getStatements :: Dis abt [Statement abt Location 'Impure] getStatements = Dis $ \_ c h -> c (statements h) h -putStatements :: [Statement abt 'Impure] -> Dis abt () +putStatements :: [Statement abt Location 'Impure] -> Dis abt () putStatements ss = - Dis $ \_ c (ListContext i _) loc -> - c () (ListContext i ss) loc + Dis $ \_ c (ListContext i _) extras -> + c () (ListContext i ss) extras ---------------------------------------------------------------- ---------------------------------------------------------------- @@ -173,7 +178,7 @@ where step :: abt '[] ('HMeasure a) - -> Statement abt 'Impure + -> Statement abt Location 'Impure -> abt '[] ('HMeasure a) step e s = #ifdef __TRACE_DISINTEGRATE__ @@ -181,13 +186,13 @@ ++ "around term " ++ show (pretty e)) $ #endif case s of - SBind x body _ -> + SBind (Location x) body _ -> -- TODO: if @body@ is dirac, then treat as 'SLet' syn (MBind :$ plugs rho (fromLazy body) :* bind x e :* End) - SLet x body _ + SLet (Location x) body _ | not (x `memberVarSet` freeVars e) -> #ifdef __TRACE_DISINTEGRATE__ - trace ("could not find location" ++ show x ++ "\n" + trace ("could not find location " ++ show x ++ "\n" ++ "in term " ++ show (pretty e) ++ "\n" ++ "given rho " ++ show (prettyAssocs rho)) $ #endif @@ -203,44 +208,46 @@ syn (Let_ :$ plugs rho (fromLazy body) :* bind x e :* End) SGuard xs pat scrutinee _ -> syn $ Case_ (plugs rho $ fromLazy scrutinee) - [ Branch pat (binds_ xs e) + [ Branch pat (binds_ (fromLocations1 xs) e) , Branch PWild (P.reject $ typeOf e) ] SWeight body _ -> syn $ Superpose_ ((plugs rho $ fromLazy body, e) :| []) ---------------------------------------------------------------- --- A location is a variable *use* instantiated at some list of indices. -data Loc :: (Hakaru -> *) -> Hakaru -> * where - Loc :: Variable a - -> [Variable 'HNat] - -> Loc ast a - MultiLoc - :: Variable a - -> [Variable 'HNat] - -> Loc ast ('HArray a) +-- An extra is a variable *use* instantiated at some list of indices. +data Extra :: (Hakaru -> *) -> Hakaru -> * where + Loc :: Location a -> [ast 'HNat] -> Extra ast a -locIndices :: Loc ast a -> [Variable 'HNat] -locIndices (Loc _ inds) = inds -locIndices (MultiLoc _ inds) = inds +extrasInds :: Extra ast a -> [ast 'HNat] +extrasInds (Loc _ inds) = inds -extendLocInds :: Variable 'HNat -> [Variable 'HNat] -> [Variable 'HNat] -extendLocInds = (:) +selectMore :: [ast 'HNat] -> ast 'HNat -> [ast 'HNat] +selectMore = flip (:) +-- Assumption: inds has no duplicates +permutes :: (ABT Term abt) + => [abt '[] 'HNat] + -> [Index (abt '[])] + -> Bool +permutes ts inds = + all isJust ts' && + length ts' == length inds && + Set.fromList (map fromJust ts') == Set.fromList (map indVar inds) + where ts' = map (\t -> caseVarSyn t Just (const Nothing)) ts + #ifdef __TRACE_DISINTEGRATE__ -prettyLoc :: Loc ast (a :: Hakaru) -> PP.Doc -prettyLoc (Loc l inds) = PP.text "Loc" PP.<+> ppVariable l - PP.<+> ppList (map ppVariable inds) -prettyLoc (MultiLoc l inds) = PP.text "MultiLoc" PP.<+> ppVariable l - PP.<+> ppList (map ppVariable inds) +prettyExtra :: (ABT Term abt) => Extra (abt '[]) a -> PP.Doc +prettyExtra (Loc (Location x) inds) = PP.text "Loc" PP.<+> ppVariable x + PP.<+> ppList (map pretty inds) -prettyLocs :: (ABT Term abt) - => Assocs (Loc (abt '[])) +prettyExtras :: (ABT Term abt) + => Assocs (Extra (abt '[])) -> PP.Doc -prettyLocs a = PP.vcat $ map go (fromAssocs a) +prettyExtras a = PP.vcat $ map go (fromAssocs a) where go (Assoc x l) = ppVariable x PP.<+> PP.text "->" PP.<+> - prettyLoc l + prettyExtra l #endif @@ -250,11 +257,12 @@ -- to make. If we turn it back into some sort of normal form, then -- it must be one preserved by 'residualizeContext'. -- --- Also, we add the list in order to support "lub" without it living in the AST. +-- Also, we add the list in order to support "lub" without it living +-- in the AST. -- TODO: really we should use LogicT... type Ans abt a = ListContext abt 'Impure - -> Assocs (Loc (abt '[])) + -> Assocs (Extra (abt '[])) -> [abt '[] ('HMeasure a)] @@ -271,11 +279,15 @@ -- -- TODO: give this a better, more informative name! -- --- N.B., This monad is used not only for both 'perform' and 'constrainOutcome', but also for 'constrainValue'. +-- N.B., This monad is used not only for both 'perform' and +-- 'constrainOutcome', but also for 'constrainValue'. newtype Dis abt x = Dis { unDis :: forall a. [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a } - -- == @Codensity (Ans abt)@, assuming 'Codensity' is poly-kinded like it should be - -- If we don't want to allow continuations that can make nondeterministic choices, then we should use the right Kan extension itself, rather than the Codensity specialization of it. + -- == @Codensity (Ans abt)@, assuming 'Codensity' is poly-kinded + -- like it should be. If we don't want to allow continuations that + -- can make nondeterministic choices, then we should use the right + -- Kan extension itself, rather than the Codensity specialization + -- of it. -- | Run a computation in the 'Dis' monad, residualizing out all the @@ -297,7 +309,11 @@ m0 [] c0 (ListContext i0 []) emptyAssocs where (Dis m0) = d >>= residualizeLocs - -- TODO: we only use dirac because 'residualizeListContext' requires it to already be a measure; unfortunately this can result in an extraneous @(>>= \x -> dirac x)@ redex at the end of the program. In principle, we should be able to eliminate that redex by changing the type of 'residualizeListContext'... + -- TODO: we only use dirac because 'residualizeListContext' + -- requires it to already be a measure; unfortunately this can + -- result in an extraneous @(>>= \x -> dirac x)@ redex at the end + -- of the program. In principle, we should be able to eliminate + -- that redex by changing the type of 'residualizeListContext'... c0 (e,rho) ss _ = [residualizeListContext ss rho (syn(Dirac :$ e :* End))] i0 = maxNextFree es @@ -359,132 +375,122 @@ -> Dis abt (abt '[] a, Assocs (abt '[])) residualizeLocs e = do ss <- getStatements - (ss', newlocs) <- foldM step ([], emptyAssocs) ss + (ss', newlocs) <- foldM step ([], emptyLAssocs) ss rho <- convertLocs newlocs putStatements (reverse ss') #ifdef __TRACE_DISINTEGRATE__ trace ("residualizeLocs: old heap:\n" ++ show (pretty_Statements ss )) $ return () trace ("residualizeLocs: new heap:\n" ++ show (pretty_Statements ss')) $ return () - locs <- getLocs - traceM ("oldlocs:\n" ++ show (prettyLocs locs) ++ "\n") + extras <- getExtras + traceM ("oldlocs:\n" ++ show (prettyExtras extras) ++ "\n") traceM ("new assoc for renaming:\n" ++ show (prettyAssocs rho)) #endif return (e, rho) where step (ss',newlocs) s = do (s',newlocs') <- residualizeLoc s - return (s':ss', insertAssocs newlocs' newlocs) + return (s':ss', insertLAssocs newlocs' newlocs) data Name (a :: Hakaru) = Name {nameHint :: Text, nameID :: Nat} -varName :: Variable a -> Name b -varName x = Name (varHint x) (varID x) +locName :: Location a -> Name b +locName (Location x) = Name (varHint x) (varID x) residualizeLoc :: (ABT Term abt) - => Statement abt 'Impure - -> Dis abt (Statement abt 'Impure, Assocs Name) + => Statement abt Location 'Impure + -> Dis abt (Statement abt Location 'Impure, LAssocs Name) residualizeLoc s = case s of SBind l _ _ -> do (s', newname) <- reifyStatement s - return (s', singletonAssocs l newname) + return (s', singletonLAssocs l newname) SLet l _ _ -> do (s', newname) <- reifyStatement s - return (s', singletonAssocs l newname) + return (s', singletonLAssocs l newname) SWeight w inds -> do - l <- freshVar Text.empty sUnit + x <- freshVar Text.empty sUnit let bodyW = Thunk $ P.weight (fromLazy w) - (s', newname) <- reifyStatement (SBind l bodyW inds) - return (s', singletonAssocs l newname) + (s', newname) <- reifyStatement (SBind (Location x) bodyW inds) + return (s', singletonLAssocs (Location x) newname) SGuard ls _ _ ixs - | null ixs -> return (s, toAssocs1 ls (fmap11 varName ls)) + | null ixs -> return (s, toLAssocs1 ls (fmap11 locName ls)) | otherwise -> error "undefined: case statement under an array" reifyStatement :: (ABT Term abt) - => Statement abt 'Impure - -> Dis abt (Statement abt 'Impure, Name a) + => Statement abt Location 'Impure + -> Dis abt (Statement abt Location 'Impure, Name a) reifyStatement s = case s of - SBind l _ [] -> return (s, varName l) + SBind l _ [] -> return (s, locName l) SBind l body (i:is) -> do let plate = Thunk . P.plateWithVar (indSize i) (indVar i) - l' <- freshVar (varHint l) (SArray (varType l)) - reifyStatement (SBind l' (plate $ fromLazy body) is) - SLet l _ [] -> return (s, varName l) + x' <- freshVar (locHint l) (SArray (locType l)) + reifyStatement (SBind (Location x') (plate $ fromLazy body) is) + SLet l _ [] -> return (s, locName l) SLet l body (i:is) -> do let array = Thunk . P.arrayWithVar (indSize i) (indVar i) - l' <- freshVar (varHint l) (SArray (varType l)) - reifyStatement (SLet l' (array $ fromLazy body) is) + x' <- freshVar (locHint l) (SArray (locType l)) + reifyStatement (SLet (Location x') (array $ fromLazy body) is) SWeight _ _ -> error "reifyStatement called on SWeight" SGuard _ _ _ _ -> error "reifyStatement called on SGuard" sizeInnermostInd :: (ABT Term abt) - => Variable (a :: Hakaru) + => Location (a :: Hakaru) -> Dis abt (abt '[] 'HNat) sizeInnermostInd l = (maybe (freeLocError l) return =<<) . select l $ \s -> do guard (length (statementInds s) >= 1) case s of - SBind l' _ ixs -> do Refl <- varEq l l' + SBind l' _ ixs -> do Refl <- locEq l l' Just $ unsafePush s >> return (indSize (head ixs)) - SLet l' _ ixs -> do Refl <- varEq l l' + SLet l' _ ixs -> do Refl <- locEq l l' Just $ unsafePush s >> return (indSize (head ixs)) SWeight _ _ -> Nothing SGuard _ _ _ _ -> error "TODO: sizeInnermostInd{SGuard}" -fromLoc :: (ABT Term abt) - => Name b - -> Sing a - -> [Variable 'HNat] - -> abt '[] a -fromLoc name typ [] = var $ Variable { varHint = nameHint name - , varID = nameID name - , varType = typ } -fromLoc name typ (i:is) = fromLoc name (SArray typ) is P.! var i +fromName :: (ABT Term abt) + => Name b + -> Sing a + -> [abt '[] 'HNat] + -> abt '[] a +fromName name typ [] = var $ Variable { varHint = nameHint name + , varID = nameID name + , varType = typ } +fromName name typ (i:is) = fromName name (SArray typ) is P.! i convertLocs :: (ABT Term abt) - => Assocs Name + => LAssocs Name -> Dis abt (Assocs (abt '[])) -convertLocs newlocs = F.foldr step emptyAssocs . fromAssocs <$> getLocs +convertLocs newlocs = F.foldr step emptyAssocs . fromAssocs <$> getExtras where build :: (ABT Term abt) - => Assoc (Loc (abt '[])) + => Assoc (Extra (abt '[])) -> Name a -> Assoc (abt '[]) - build (Assoc x loc) name = - Assoc x (fromLoc name (varType x) - (case loc of Loc _ js -> js; MultiLoc _ js -> js)) - step assoc@(Assoc _ loc) = insertAssoc $ - case loc of + build (Assoc x extra) name = + Assoc x (fromName name (varType x) (extrasInds extra)) + step assoc@(Assoc _ extra) = insertAssoc $ + case extra of Loc l _ -> maybe (freeLocError l) (build assoc) - (lookupAssoc l newlocs) - MultiLoc l _ -> maybe (freeLocError l) - (build assoc) - (lookupAssoc l newlocs) + (lookupLAssoc l newlocs) -freeLocError :: Variable (a :: Hakaru) -> b +freeLocError :: Location (a :: Hakaru) -> b freeLocError l = error $ "Found a free location " ++ show l +zipInds :: (ABT Term abt) + => [Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[]) +zipInds inds ts + | length inds /= length ts + = error "zipInds: argument lists must have the same length" + | otherwise = toAssocs $ zipWith Assoc (map indVar inds) ts + apply :: (ABT Term abt) - => [(Index (abt '[]), Index (abt '[]))] + => [Index (abt '[])] + -> [Index (abt '[])] -> abt '[] a -> Dis abt (abt '[] a) -apply ijs e = do locs <- fromAssocs <$> getLocs - rho' <- foldM step rho locs - return (renames rho' e) - where - rho = toAssocs $ map (\(i,j) -> Assoc (indVar i) (indVar j)) ijs - step r (Assoc x loc) = - let inds = locIndices loc - check i = lookupAssoc i rho - inds' = map (\i -> fromMaybe i (check i)) inds - in if (any isJust (map check inds)) - then do x' <- case loc of - Loc l _ -> mkLoc Text.empty l inds' - MultiLoc l _ -> mkMultiLoc Text.empty l inds' - return (insertAssoc (Assoc x x') r) - else return r +apply is js e = extSubsts (zipInds is (map fromIndex js)) e extendIndices :: (ABT Term abt) @@ -500,7 +506,7 @@ = j : js -- give better name -statementInds :: Statement abt p -> [Index (abt '[])] +statementInds :: Statement abt Location p -> [Index (abt '[])] statementInds (SBind _ _ i) = i statementInds (SLet _ _ i) = i statementInds (SWeight _ i) = i @@ -508,62 +514,51 @@ statementInds (SStuff0 _ i) = i statementInds (SStuff1 _ _ i) = i -getLocs :: (ABT Term abt) - => Dis abt (Assocs (Loc (abt '[]))) -getLocs = Dis $ \_ c h l -> c l h l - -putLocs :: (ABT Term abt) - => Assocs (Loc (abt '[])) - -> Dis abt () -putLocs l = Dis $ \_ c h _ -> c () h l +getExtras :: (ABT Term abt) + => Dis abt (Assocs (Extra (abt '[]))) +getExtras = Dis $ \_ c h l -> c l h l -insertLoc :: (ABT Term abt) - => Variable a - -> Loc (abt '[]) a +putExtras :: (ABT Term abt) + => Assocs (Extra (abt '[])) -> Dis abt () -insertLoc v loc = +putExtras l = Dis $ \_ c h _ -> c () h l + +insertExtra :: (ABT Term abt) + => Variable a + -> Extra (abt '[]) a + -> Dis abt () +insertExtra v extra = Dis $ \_ c h l -> c () h $ - insertAssoc (Assoc v loc) l + insertAssoc (Assoc v extra) l -adjustLoc :: (ABT Term abt) - => Variable (a :: Hakaru) - -> (Assoc (Loc (abt '[])) -> Assoc (Loc (abt '[]))) - -> Dis abt () -adjustLoc x f = do - locs <- getLocs - putLocs $ adjustAssoc x f locs +adjustExtra :: (ABT Term abt) + => Variable (a :: Hakaru) + -> (Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[]))) + -> Dis abt () +adjustExtra x f = do + extras <- getExtras + putExtras $ adjustAssoc x f extras mkLoc :: (ABT Term abt) => Text - -> Variable (a :: Hakaru) - -> [Variable 'HNat] + -> Location (a :: Hakaru) + -> [abt '[] 'HNat] -> Dis abt (Variable a) -mkLoc hint s inds = do - x <- freshVar hint (varType s) - insertLoc x (Loc s inds) +mkLoc hint l inds = do + x <- freshVar hint (locType l) + insertExtra x (Loc l inds) return x mkLocs :: (ABT Term abt) - => List1 Variable (xs :: [Hakaru]) - -> [Variable 'HNat] + => List1 Location (xs :: [Hakaru]) + -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs) mkLocs Nil1 _ = return Nil1 -mkLocs (Cons1 x xs) inds = Cons1 - <$> mkLoc Text.empty x inds - <*> mkLocs xs inds - -mkMultiLoc - :: (ABT Term abt) - => Text - -> Variable a - -> [Variable 'HNat] - -> Dis abt (Variable ('HArray a)) -mkMultiLoc hint s inds = do - x' <- freshVar hint (SArray $ varType s) - insertLoc x' (MultiLoc s inds) - return x' +mkLocs (Cons1 l ls) inds = Cons1 + <$> mkLoc Text.empty l inds + <*> mkLocs ls inds instance Functor (Dis abt) where fmap f (Dis m) = Dis $ \i c -> m i (c . f) @@ -589,21 +584,26 @@ Dis $ \_ c (ListContext n ss) -> c n (ListContext (n+1) ss) - freshenStatement s = + -- Note: we assume that freshLocStatement is never called on a + -- statement already on the heap (list context) + freshLocStatement s = case s of - SWeight _ _ -> return (s, mempty) + SWeight w e -> return (SWeight w e, mempty) SBind x body i -> do - l <- freshenVar x - x' <- mkLoc (varHint x) l (map indVar i) - return (SBind l body i, singletonAssocs x x') + x' <- freshenVar x + let l = Location x' + v <- mkLoc (locHint l) l (map fromIndex i) + return (SBind l body i, singletonAssocs x v) SLet x body i -> do - l <- freshenVar x - x' <- mkLoc (varHint x) l (map indVar i) - return (SLet l body i, singletonAssocs x x') + x' <- freshenVar x + let l = Location x' + v <- mkLoc (locHint l) l (map fromIndex i) + return (SLet l body i, singletonAssocs x v) SGuard xs pat scrutinee i -> do - ls <- freshenVars xs - xs' <- mkLocs ls (map indVar i) - return (SGuard ls pat scrutinee i, toAssocs1 xs xs') + xs' <- freshenVars xs + let ls = locations1 xs' + vs <- mkLocs ls (map fromIndex i) + return (SGuard ls pat scrutinee i, toAssocs1 xs vs) getIndices = Dis $ \i c -> c i @@ -637,31 +637,127 @@ unsafePushes ss return (Just r) + substVar x e z = do + extras <- getExtras + let defaultResult = return (var z) + case (lookupAssoc z extras) of + Nothing -> defaultResult + Just (Loc l inds) -> + if any (memberVarSet x . freeVars) inds + then do inds' <- mapM (extSubst x e) inds + var <$> mkLoc Text.empty l inds' + else defaultResult + + extFreeVars e = do + extras <- getExtras + let fvs1 = freeVars e + indFVs (SomeVariable v) = + case (lookupAssoc v extras) of + Nothing -> emptyVarSet + Just (Loc _ is) -> foldr (unionVarSet.freeVars) emptyVarSet is + locVars (SomeVariable v) b = + case (lookupAssoc v extras) of + Nothing -> b + Just (Loc l _) -> insertVarSet (fromLocation l) b + fvs2 = foldr (unionVarSet.indFVs) fvs1 (fromVarSet fvs1) + return $ foldr locVars emptyVarSet (fromVarSet fvs2) + + -- | The forward disintegrator's function for evaluating case + -- expressions. First we try calling 'defaultCaseEvaluator' which + -- will evaluate the scrutinee and select the matching branch (if + -- any). But that doesn't work out in general, since the scrutinee + -- may contain heap-bound variables. So our fallback definition + -- will push a 'SGuard' onto the heap and then continue evaluating + -- each branch (thereby duplicating the continuation, calling it + -- once on each branch). + evaluateCase evaluate_ = evaluateCase_ + where + evaluateCase_ :: CaseEvaluator abt (Dis abt) + evaluateCase_ e bs = + defaultCaseEvaluator evaluate_ e bs + <|> evaluateBranches e bs + + evaluateBranches :: CaseEvaluator abt (Dis abt) + evaluateBranches e = choose . map evaluateBranch + where + evaluateBranch (Branch pat body) = + let (vars,body') = caseBinds body + in getIndices >>= \i -> + push (SGuard vars pat (Thunk e) i) body' + >>= evaluate_ + + evaluateVar perform evaluate_ x = + do extras <- getExtras + -- If we get 'Nothing', then it turns out @x@ is a free variable + maybe (return $ Neutral (var x)) lookForLoc (lookupAssoc x extras) + where lookForLoc (Loc l jxs) = + (maybe (freeLocError l) return =<<) . select l $ \s -> + case s of + SBind l' e ixs -> do + Refl <- locEq l l' + Just $ do + w <- withIndices ixs $ perform (caseLazy e fromWhnf id) + unsafePush (SLet l (Whnf_ w) ixs) +#ifdef __TRACE_DISINTEGRATE__ + trace ("-- updated " + ++ show (ppStatement 11 s) + ++ " to " + ++ show (ppStatement 11 (SLet l (Whnf_ w) ixs)) + ) $ return () +#endif + extSubsts (zipInds ixs jxs) (fromWhnf w) >>= evaluate_ + SLet l' e ixs -> do + Refl <- locEq l l' + Just $ do + w <- withIndices ixs $ caseLazy e return evaluate_ + unsafePush (SLet l (Whnf_ w) ixs) + extSubsts (zipInds ixs jxs) (fromWhnf w) >>= evaluate_ + -- This does not bind any variables, so it definitely can't match. + SWeight _ _ -> Nothing + -- This does bind variables, + -- but there's no expression we can return for it + -- because the variables are untouchable\/abstract. + SGuard ls pat scrutinee i -> Just . return . Neutral $ var x + + withIndices :: [Index (abt '[])] -> Dis abt a -> Dis abt a withIndices inds (Dis m) = Dis $ \_ c -> m inds c -- | Not exported because we only need it for defining 'select' on 'Dis'. -unsafePop :: Dis abt (Maybe (Statement abt 'Impure)) +unsafePop :: Dis abt (Maybe (Statement abt Location 'Impure)) unsafePop = - Dis $ \_ c h@(ListContext i ss) loc -> + Dis $ \_ c h@(ListContext i ss) extras -> case ss of - [] -> c Nothing h loc - s:ss' -> c (Just s) (ListContext i ss') loc + [] -> c Nothing h extras + s:ss' -> c (Just s) (ListContext i ss') extras pushPlate :: (ABT Term abt) => abt '[] 'HNat -> abt '[ 'HNat ] ('HMeasure a) - -> Dis abt (Variable ('HArray a)) + -> Dis abt (abt '[] ('HArray a)) pushPlate n e = caseBind e $ \x body -> do inds <- getIndices i <- freshInd n - p <- freshVar Text.empty (sUnMeasure $ typeOf body) - unsafePush (SBind p (Thunk $ rename x (indVar i) body) - (extendIndices i inds)) - mkMultiLoc Text.empty p (map indVar inds) + p <- Location <$> freshVar Text.empty (sUnMeasure $ typeOf body) + let inds' = extendIndices i inds + unsafePush (SBind p (Thunk $ rename x (indVar i) body) inds') + v <- mkLoc Text.empty p $ map fromIndex inds' + return $ P.arrayWithVar n (indVar i) (var v) +-- | Calls unsafePush +pushWeight :: (ABT Term abt) => abt '[] 'HProb -> Dis abt () +pushWeight w = do + inds <- getIndices + unsafePush $ SWeight (Thunk w) inds + +-- | Calls unsafePush +pushGuard :: (ABT Term abt) => abt '[] HBool -> Dis abt () +pushGuard b = do + inds <- getIndices + unsafePush $ SGuard Nil1 pTrue (Thunk b) inds + ---------------------------------------------------------------- ---------------------------------------------------------------- @@ -697,7 +793,6 @@ x <- freshVar hint typ Dis $ \_ c h l -> (f . bind x) <$> c x h l - -- This function was called @lift@ in the finally-tagless code. -- | Emit an 'MBind' (i.e., \"@m >>= \x ->@\") and return the -- variable thus bound (i.e., @x@). @@ -706,6 +801,17 @@ emit Text.empty (sUnMeasure $ typeOf m) $ \e -> syn (MBind :$ m :* e :* End) +emitMBind2 :: (ABT Term abt) => abt '[] ('HMeasure a) -> Dis abt (abt '[] a) +emitMBind2 m = do + inds <- getIndices + let b = Whnf_ $ fromMaybe (error "emitMBind2: non-hnf term") (toWhnf m) + typ = sUnMeasure $ typeOf m + l <- Location <$> freshVar Text.empty typ + (SBind l' b' _, name) <- reifyStatement (SBind l b inds) + let (idx, p) = (fromName name typ (map fromIndex inds), fromLazy b') + Dis $ \_ c h ex -> + c idx h ex >>= \e -> + return $ syn (MBind :$ p :* bind (fromLocation l') e :* End) -- | A smart constructor for emitting let-bindings. If the input -- is already a variable then we just return it; otherwise we emit @@ -819,7 +925,8 @@ emitMBind_ m = emit_ (m P.>>) --- TODO: if the argument is a value, then we can evaluate the 'P.if_' immediately rather than emitting it. +-- TODO: if the argument is a value, then we can evaluate the 'P.if_' +-- immediately rather than emitting it. -- | Emit an assertion that the condition is true. emitGuard :: (ABT Term abt) => abt '[] HBool -> Dis abt () emitGuard b = emit_ (P.withGuard b) -- == emit_ $ \m -> P.if_ b m P.reject diff --git a/haskell/Language/Hakaru/Evaluation/EvalMonad.hs b/haskell/Language/Hakaru/Evaluation/EvalMonad.hs --- a/haskell/Language/Hakaru/Evaluation/EvalMonad.hs +++ b/haskell/Language/Hakaru/Evaluation/EvalMonad.hs @@ -47,7 +47,7 @@ import Language.Hakaru.Syntax.DatumABT import Language.Hakaru.Syntax.AST import Language.Hakaru.Evaluation.Types -import Language.Hakaru.Evaluation.Lazy (TermEvaluator, evaluate, defaultCaseEvaluator) +import Language.Hakaru.Evaluation.Lazy (evaluate) import Language.Hakaru.Evaluation.PEvalMonad (ListContext(..)) @@ -83,7 +83,7 @@ -- -- | Call 'evaluate' on a term. This variant returns something in the 'Eval' monad so you can string multiple evaluation calls together. For the non-monadic version, see 'runPureEvaluate'. pureEvaluate :: (ABT Term abt) => TermEvaluator abt (Eval abt) -pureEvaluate = evaluate (brokenInvariant "perform") defaultCaseEvaluator +pureEvaluate = evaluate (brokenInvariant "perform") ---------------------------------------------------------------- @@ -124,10 +124,10 @@ foldl step e0 . statements where -- TODO: make paremetric in the purity, so we can combine 'residualizeListContext' with this function. - step :: abt '[] a -> Statement abt 'Pure -> abt '[] a + step :: abt '[] a -> Statement abt Location 'Pure -> abt '[] a step e s = case s of - SLet x body _ + SLet (Location x) body _ | not (x `memberVarSet` freeVars e) -> e -- TODO: if used exactly once in @e@, then inline. | otherwise -> @@ -187,9 +187,9 @@ unsafePushes ss return (Just r) --- TODO: make paremetric in the purity +-- TODO: make parametric in the purity -- | Not exported because we only need it for defining 'select' on 'Eval'. -unsafePop :: Eval abt (Maybe (Statement abt 'Pure)) +unsafePop :: Eval abt (Maybe (Statement abt Location 'Pure)) unsafePop = Eval $ \c h@(ListContext i ss) -> case ss of diff --git a/haskell/Language/Hakaru/Evaluation/ExpectMonad.hs b/haskell/Language/Hakaru/Evaluation/ExpectMonad.hs --- a/haskell/Language/Hakaru/Evaluation/ExpectMonad.hs +++ b/haskell/Language/Hakaru/Evaluation/ExpectMonad.hs @@ -49,7 +49,7 @@ import Language.Hakaru.Syntax.Variable (memberVarSet) import Language.Hakaru.Syntax.AST hiding (Expect) import Language.Hakaru.Evaluation.Types -import Language.Hakaru.Evaluation.Lazy (TermEvaluator, evaluate, defaultCaseEvaluator) +import Language.Hakaru.Evaluation.Lazy (evaluate) import Language.Hakaru.Evaluation.PEvalMonad (ListContext(..)) @@ -79,10 +79,10 @@ foldl step e0 . statements where -- TODO: make paremetric in the purity, so we can combine 'residualizeListContext' with this function. - step :: abt '[] 'HProb -> Statement abt 'ExpectP -> abt '[] 'HProb + step :: abt '[] 'HProb -> Statement abt Location 'ExpectP -> abt '[] 'HProb step e s = case s of - SLet x body _ + SLet (Location x) body _ -- BUG: this trick for dropping unused let-bindings doesn't seem to work anymore... (cf., 'Tests.Expect.test4') | not (x `memberVarSet` freeVars e) -> e -- TODO: if used exactly once in @e@, then inline. @@ -99,7 +99,7 @@ pureEvaluate :: (ABT Term abt) => TermEvaluator abt (Expect abt) -pureEvaluate = evaluate (brokenInvariant "perform") defaultCaseEvaluator +pureEvaluate = evaluate (brokenInvariant "perform") brokenInvariant :: String -> a brokenInvariant loc = error (loc ++ ": Expect's invariant broken") @@ -185,7 +185,7 @@ -- TODO: make paremetric in the purity -- | Not exported because we only need it for defining 'select' on 'Expect'. -unsafePop :: Expect abt (Maybe (Statement abt 'ExpectP)) +unsafePop :: Expect abt (Maybe (Statement abt Location 'ExpectP)) unsafePop = Expect $ \c h@(ListContext i ss) -> case ss of diff --git a/haskell/Language/Hakaru/Evaluation/Lazy.hs b/haskell/Language/Hakaru/Evaluation/Lazy.hs --- a/haskell/Language/Hakaru/Evaluation/Lazy.hs +++ b/haskell/Language/Hakaru/Evaluation/Lazy.hs @@ -29,18 +29,11 @@ -- like my old one. ---------------------------------------------------------------- module Language.Hakaru.Evaluation.Lazy - ( - -- * Lazy partial evaluation - TermEvaluator - , MeasureEvaluator - , CaseEvaluator - , VariableEvaluator - , evaluate + ( evaluate -- ** Helper functions - , update - , defaultCaseEvaluator - , toStatements - + , evaluateNaryOp + , evaluatePrimOp + , evaluateArrayOp -- ** Helpers that should really go away , Interp(..), reifyPair ) where @@ -91,44 +84,16 @@ -- to select out subtypes of the generic versions. --- | A function for evaluating any term to weak-head normal form. -type TermEvaluator abt m = - forall a. abt '[] a -> m (Whnf abt a) - --- | A function for \"performing\" an 'HMeasure' monadic action. --- This could mean actual random sampling, or simulated sampling --- by generating a new term and returning the newly bound variable, --- or anything else. -type MeasureEvaluator abt m = - forall a. abt '[] ('HMeasure a) -> m (Whnf abt a) - --- | A function for evaluating any case-expression to weak-head --- normal form. -type CaseEvaluator abt m = - forall a b. abt '[] a -> [Branch a abt b] -> m (Whnf abt b) - --- | A function for evaluating any variable to weak-head normal form. -type VariableEvaluator abt m = - forall a. Variable a -> m (Whnf abt a) - - -- | Lazy partial evaluation with some given \"perform\" and --- \"evaluateCase\" functions. The first argument to @evaluateCase@ --- will be the 'TermEvaluator' we're constructing (thus tying the --- knot). N.B., if @p ~ 'Pure@ then the \"perform\" function will --- never be called. --- --- We factor out the 'CaseEvaluator' because some code (e.g., --- disintegration) may need to do something special rather than --- just relying on the 'defaultCaseEvaluator' implementation. +-- \"evaluateCase\" functions. N.B., if @p ~ 'Pure@ then the +-- \"perform\" function will never be called. evaluate :: forall abt m p . (ABT Term abt, EvaluationMonad abt m p) => MeasureEvaluator abt m - -> (TermEvaluator abt m -> CaseEvaluator abt m) -> TermEvaluator abt m {-# INLINE evaluate #-} -evaluate perform evaluateCase = evaluate_ +evaluate perform = evaluate_ where evaluateCase_ :: CaseEvaluator abt m evaluateCase_ = evaluateCase evaluate_ @@ -138,13 +103,14 @@ #ifdef __TRACE_DISINTEGRATE__ trace ("-- evaluate_: " ++ show (pretty e0)) $ #endif - caseVarSyn e0 (update perform evaluate_) $ \t -> + caseVarSyn e0 (evaluateVar perform evaluate_) $ \t -> case t of -- Things which are already WHNFs Literal_ v -> return . Head_ $ WLiteral v Datum_ d -> return . Head_ $ WDatum d Empty_ typ -> return . Head_ $ WEmpty typ Array_ e1 e2 -> return . Head_ $ WArray e1 e2 + ArrayLiteral_ es -> return . Head_ $ WArrayLiteral es Lam_ :$ e1 :* End -> return . Head_ $ WLam e1 Dirac :$ e1 :* End -> return . Head_ $ WDirac e1 MBind :$ e1 :* e2 :* End -> return . Head_ $ WMBind e1 e2 @@ -172,13 +138,13 @@ -- call-by-name: caseBind f $ \x f' -> do i <- getIndices - push (SLet x (Thunk e2) i) f' evaluate_ + push (SLet x (Thunk e2) i) f' >>= evaluate_ evaluateApp _ = error "evaluate{App_}: the impossible happened" Let_ :$ e1 :* e2 :* End -> do i <- getIndices caseBind e2 $ \x e2' -> - push (SLet x (Thunk e1) i) e2' evaluate_ + push (SLet x (Thunk e1) i) e2' >>= evaluate_ CoerceTo_ c :$ e1 :* End -> coerceTo c <$> evaluate_ e1 UnsafeFrom_ c :$ e1 :* End -> coerceFrom c <$> evaluate_ e1 @@ -203,104 +169,7 @@ _ :$ _ -> error "evaluate: the impossible happened" --- | A simple 'CaseEvaluator' which uses the 'DatumEvaluator' to --- force the scrutinee, and if 'matchBranches' succeeds then we --- call the 'TermEvaluator' to continue evaluating the body of the --- matched branch. If we 'GotStuck' then we return a 'Neutral' term --- of the case expression itself (n.b, any side effects from having --- called the 'DatumEvaluator' will still persist when returning --- this neutral term). If we didn't get stuck and yet none of the --- branches matches, then we throw an exception. -defaultCaseEvaluator - :: forall abt m p - . (ABT Term abt, EvaluationMonad abt m p) - => TermEvaluator abt m - -> CaseEvaluator abt m -{-# INLINE defaultCaseEvaluator #-} -defaultCaseEvaluator evaluate_ = evaluateCase_ - where - -- TODO: At present, whenever we residualize a case expression we'll - -- generate a 'Neutral' term which will, when run, repeat the work - -- we're doing in the evaluation here. We could eliminate this - -- redundancy by introducing a new variable for @e@ each time this - -- function is called--- if only we had some way of getting those - -- variables put into the right place for when we residualize the - -- original scrutinee... - -- - -- N.B., 'DatumEvaluator' is a rank-2 type so it requires a signature - evaluateDatum :: DatumEvaluator (abt '[]) m - evaluateDatum e = viewWhnfDatum <$> evaluate_ e - - evaluateCase_ :: CaseEvaluator abt m - evaluateCase_ e bs = do - match <- matchBranches evaluateDatum e bs - case match of - Nothing -> - -- TODO: print more info about where this error - -- happened - -- - -- TODO: rather than throwing a Haskell error, - -- instead capture the possibility of failure in - -- the 'EvaluationMonad' monad. - error "defaultCaseEvaluator: non-exhaustive patterns in case!" - Just GotStuck -> - return . Neutral . syn $ Case_ e bs - Just (Matched ss body) -> - pushes (toStatements ss) body evaluate_ - - -toStatements :: Assocs (abt '[]) -> [Statement abt p] -toStatements = map (\(Assoc x e) -> SLet x (Thunk e) []) . fromAssocs - - ---------------------------------------------------------------- --- TODO: figure out how to abstract this so it can be reused by --- 'constrainValue'. Especially the 'SBranch case of 'step' --- --- TODO: we could speed up the case for free variables by having --- the 'Context' also keep track of the largest free var. That way, --- we can just check up front whether @varID x < nextFreeVarID@. --- Of course, we'd have to make sure we've sufficiently renamed all --- bound variables to be above @nextFreeVarID@; but then we have to --- do that anyways. -update - :: forall abt m p - . (ABT Term abt, EvaluationMonad abt m p) - => MeasureEvaluator abt m - -> TermEvaluator abt m - -> VariableEvaluator abt m -update perform evaluate_ = \x -> - -- If we get 'Nothing', then it turns out @x@ is a free variable - fmap (maybe (Neutral $ var x) id) . select x $ \s -> - case s of - SBind y e i -> do - Refl <- varEq x y - Just $ do - w <- perform $ caseLazy e fromWhnf id - unsafePush (SLet x (Whnf_ w) i) -#ifdef __TRACE_DISINTEGRATE__ - trace ("-- updated " - ++ show (ppStatement 11 s) - ++ " to " - ++ show (ppStatement 11 (SLet x (Whnf_ w) i)) - ) $ return () -#endif - return w - SLet y e i -> do - Refl <- varEq x y - Just $ do - w <- caseLazy e return evaluate_ - unsafePush (SLet x (Whnf_ w) i) - return w - -- These two don't bind any variables, so they definitely can't match. - SWeight _ _ -> Nothing - SStuff0 _ _ -> Nothing - -- These two do bind variables, but there's no expression we can return for them because the variables are untouchable\/abstract. - SStuff1 _ _ _ -> Just . return . Neutral $ var x - SGuard ys pat scrutinee i -> Just . return . Neutral $ var x - - ----------------------------------------------------------------- -- BUG: need to improve the types so they can capture polymorphic data types -- BUG: this is a **really gross** hack. If we can avoid it, we should!!! class Interp a a' | a -> a' where @@ -546,37 +415,37 @@ evaluateArrayOp evaluate_ = go where go o@(Index _) = \(e1 :* e2 :* End) -> do + let -- idxCode :: abt '[] ('HArray a) -> abt '[] 'HNat -> abt '[] a + -- idxCode a i = Neutral $ syn (ArrayOp_ o :$ a :* i :* End) w1 <- evaluate_ e1 case w1 of - Neutral e1' -> + Neutral e1' -> return . Neutral $ syn (ArrayOp_ o :$ e1' :* e2 :* End) - Head_ v1 -> - error "TODO: evaluateArrayOp{Index}{Head_}" + Head_ (WArray _ b) -> + caseBind b $ \x body -> extSubst x e2 body >>= evaluate_ + Head_ (WEmpty _) -> + error "TODO: evaluateArrayOp{Index}{Head_ (WEmpty _)}" + Head_ (WArrayLiteral arr) -> + do w2 <- evaluate_ e2 + case w2 of + Head_ (WLiteral (LNat n)) -> return . Neutral $ + arr !! fromInteger (fromNatural n) + _ -> return . Neutral $ + syn (ArrayOp_ o :$ fromWhnf w1 :* fromWhnf w2 :* End) + _ -> error "evaluateArrayOp{Index}: uknown whnf of array type" go o@(Size _) = \(e1 :* End) -> do w1 <- evaluate_ e1 case w1 of Neutral e1' -> return . Neutral $ syn (ArrayOp_ o :$ e1' :* End) - Head_ v1 -> - case head2array v1 of - WAEmpty -> return . Head_ $ WLiteral (LNat 0) - WAArray e3 _ -> evaluate_ e3 + Head_ (WEmpty _) -> return . Head_ $ WLiteral (LNat 0) + Head_ (WArray e2 _) -> evaluate_ e2 + Head_ (WArrayLiteral es) -> return . Head_ . WLiteral . + primCoerceFrom (Signed HRing_Int) . + LInt . toInteger $ length es go (Reduce _) = \(e1 :* e2 :* e3 :* End) -> error "TODO: evaluateArrayOp{Reduce}" - - -data ArrayHead :: ([Hakaru] -> Hakaru -> *) -> Hakaru -> * where - WAEmpty :: ArrayHead abt a - WAArray - :: !(abt '[] 'HNat) - -> !(abt '[ 'HNat] a) - -> ArrayHead abt a - -head2array :: Head abt ('HArray a) -> ArrayHead abt a -head2array (WEmpty _) = WAEmpty -head2array (WArray e1 e2) = WAArray e1 e2 - ---------------------------------------------------------------- -- TODO: maybe we should adjust 'Whnf' to have a third option for diff --git a/haskell/Language/Hakaru/Evaluation/PEvalMonad.hs b/haskell/Language/Hakaru/Evaluation/PEvalMonad.hs --- a/haskell/Language/Hakaru/Evaluation/PEvalMonad.hs +++ b/haskell/Language/Hakaru/Evaluation/PEvalMonad.hs @@ -117,7 +117,7 @@ data ListContext (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity) = ListContext { nextFreshNat :: {-# UNPACK #-} !Nat - , statements :: [Statement abt p] + , statements :: [Statement abt Location p] } @@ -178,17 +178,17 @@ \ss e0 -> foldl (flip step) e0 (statements ss) where step - :: Statement abt p + :: Statement abt Location p -> P p abt '[] a -> P p abt '[] a - step (SLet x body _) = mapP $ residualizeLet x body - step (SBind x body _) = mapPImpure $ \e -> + step (SLet (Location x) body _) = mapP $ residualizeLet x body + step (SBind (Location x) body _) = mapPImpure $ \e -> -- TODO: if @body@ is dirac, then treat as 'SLet' syn (MBind :$ fromLazy body :* bind x e :* End) step (SGuard xs pat scrutinee _) = mapPImpure $ \e -> -- TODO: avoid adding the 'PWild' branch if we know @pat@ covers the type syn $ Case_ (fromLazy scrutinee) - [ Branch pat $ binds_ xs e + [ Branch pat $ binds_ (fromLocations1 xs) e , Branch PWild $ P.reject (typeOf e) ] step (SWeight body _) = mapPImpure $ P.withWeight (fromLazy body) @@ -355,7 +355,7 @@ return (Just r) -- | Not exported because we only need it for defining 'select' on 'PEval'. -unsafePop :: PEval abt p m (Maybe (Statement abt p)) +unsafePop :: PEval abt p m (Maybe (Statement abt Location p)) unsafePop = PEval $ \c h@(ListContext i ss) -> case ss of diff --git a/haskell/Language/Hakaru/Evaluation/Types.hs b/haskell/Language/Hakaru/Evaluation/Types.hs --- a/haskell/Language/Hakaru/Evaluation/Types.hs +++ b/haskell/Language/Hakaru/Evaluation/Types.hs @@ -2,6 +2,7 @@ , GADTs , KindSignatures , DataKinds + , PolyKinds , TypeOperators , Rank2Types , BangPatterns @@ -11,6 +12,7 @@ , FlexibleInstances , UndecidableInstances , EmptyCase + , ScopedTypeVariables #-} {-# OPTIONS_GHC -Wall -fwarn-tabs #-} @@ -42,9 +44,19 @@ , getLazyVariable, isLazyVariable , getLazyLiteral, isLazyLiteral + -- * Lazy partial evaluation + , TermEvaluator + , MeasureEvaluator + , CaseEvaluator + , VariableEvaluator + -- * The monad for partial evaluation , Purity(..), Statement(..), statementVars, isBoundBy - , Index, indVar, indSize + , Index, indVar, indSize, fromIndex + , Location(..), locEq, locHint, locType, locations1 + , fromLocation, fromLocations1, freshenLoc, freshenLocs + , LAssoc, LAssocs , emptyLAssocs, singletonLAssocs + , toLAssocs1, insertLAssocs, lookupLAssoc #ifdef __TRACE_DISINTEGRATE__ , ppList , ppInds @@ -54,13 +66,17 @@ , prettyAssocs #endif , EvaluationMonad(..) + , defaultCaseEvaluator + , toVarStatements + , extSubst + , extSubsts , freshVar , freshenVar , Hint(..), freshVars , freshenVars , freshInd {- TODO: should we expose these? - , freshenStatement + , freshLocStatement , push_ -} , push @@ -73,6 +89,7 @@ import Data.Monoid (Monoid(..)) import Data.Functor ((<$>)) import Control.Applicative (Applicative(..)) +import Data.Traversable #endif import Control.Arrow ((***)) import qualified Data.Foldable as F @@ -88,6 +105,9 @@ import Language.Hakaru.Types.Coercion import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.DatumCase (DatumEvaluator, + MatchResult(..), + matchBranches) import Language.Hakaru.Syntax.AST.Eq (alphaEq) -- import Language.Hakaru.Syntax.TypeOf import Language.Hakaru.Syntax.ABT @@ -96,6 +116,7 @@ #ifdef __TRACE_DISINTEGRATE__ import qualified Text.PrettyPrint as PP import Language.Hakaru.Pretty.Haskell +import Debug.Trace (trace) #endif ---------------------------------------------------------------- @@ -134,6 +155,8 @@ WDatum :: !(Datum (abt '[]) (HData' t)) -> Head abt (HData' t) WEmpty :: !(Sing ('HArray a)) -> Head abt ('HArray a) WArray :: !(abt '[] 'HNat) -> !(abt '[ 'HNat] a) -> Head abt ('HArray a) + WArrayLiteral + :: [abt '[] a] -> Head abt ('HArray a) WLam :: !(abt '[ a ] b) -> Head abt (a ':-> b) -- Measure heads (N.B., not simply @abt '[] ('HMeasure _)@) @@ -195,6 +218,7 @@ fromHead (WDatum d) = syn (Datum_ d) fromHead (WEmpty typ) = syn (Empty_ typ) fromHead (WArray e1 e2) = syn (Array_ e1 e2) +fromHead (WArrayLiteral es) = syn (ArrayLiteral_ es) fromHead (WLam e1) = syn (Lam_ :$ e1 :* End) fromHead (WMeasureOp o es) = syn (MeasureOp_ o :$ es) fromHead (WDirac e1) = syn (Dirac :$ e1 :* End) @@ -218,6 +242,7 @@ Datum_ d -> Just $ WDatum d Empty_ typ -> Just $ WEmpty typ Array_ e1 e2 -> Just $ WArray e1 e2 + ArrayLiteral_ es -> Just $ WArrayLiteral es Lam_ :$ e1 :* End -> Just $ WLam e1 MeasureOp_ o :$ es -> Just $ WMeasureOp o es Dirac :$ e1 :* End -> Just $ WDirac e1 @@ -236,6 +261,7 @@ fmap21 f (WDatum d) = WDatum (fmap11 f d) fmap21 _ (WEmpty typ) = WEmpty typ fmap21 f (WArray e1 e2) = WArray (f e1) (f e2) + fmap21 f (WArrayLiteral es) = WArrayLiteral (fmap f es) fmap21 f (WLam e1) = WLam (f e1) fmap21 f (WMeasureOp o es) = WMeasureOp o (fmap21 f es) fmap21 f (WDirac e1) = WDirac (f e1) @@ -254,6 +280,7 @@ foldMap21 f (WDatum d) = foldMap11 f d foldMap21 _ (WEmpty _) = mempty foldMap21 f (WArray e1 e2) = f e1 `mappend` f e2 + foldMap21 f (WArrayLiteral es) = F.foldMap f es foldMap21 f (WLam e1) = f e1 foldMap21 f (WMeasureOp _ es) = foldMap21 f es foldMap21 f (WDirac e1) = f e1 @@ -272,6 +299,7 @@ traverse21 f (WDatum d) = WDatum <$> traverse11 f d traverse21 _ (WEmpty typ) = pure $ WEmpty typ traverse21 f (WArray e1 e2) = WArray <$> f e1 <*> f e2 + traverse21 f (WArrayLiteral es) = WArrayLiteral <$> traverse f es traverse21 f (WLam e1) = WLam <$> f e1 traverse21 f (WMeasureOp o es) = WMeasureOp o <$> traverse21 f es traverse21 f (WDirac e1) = WDirac <$> f e1 @@ -300,7 +328,10 @@ -- TODO: would it be helpful to track which variable it's blocked -- on? To do so we'd need 'GotStuck' to return that info... -- - -- TODO: is there some /clean/ way to ensure that the neutral term is exactly a chain of blocked redexes? That is, we want to be able to pull out neutral 'Case_' terms; so we want to make sure they're not wrapped in let-bindings, coercions, etc. + -- TODO: is there some /clean/ way to ensure that the neutral term + -- is exactly a chain of blocked redexes? That is, we want to be + -- able to pull out neutral 'Case_' terms; so we want to make sure + -- they're not wrapped in let-bindings, coercions, etc. -- | Forget that something is a WHNF. fromWhnf :: (ABT Term abt) => Whnf abt a -> abt '[] a @@ -319,19 +350,6 @@ caseWhnf (Neutral e) _ k = k e -instance Functor21 Whnf where - fmap21 f (Head_ v) = Head_ (fmap21 f v) - fmap21 f (Neutral e) = Neutral (f e) - -instance Foldable21 Whnf where - foldMap21 f (Head_ v) = foldMap21 f v - foldMap21 f (Neutral e) = f e - -instance Traversable21 Whnf where - traverse21 f (Head_ v) = Head_ <$> traverse21 f v - traverse21 f (Neutral e) = Neutral <$> f e - - -- | Given some WHNF, try to extract a 'Datum' from it. viewWhnfDatum :: (ABT Term abt) @@ -429,19 +447,6 @@ caseLazy (Whnf_ e) k _ = k e caseLazy (Thunk e) _ k = k e -instance Functor21 Lazy where - fmap21 f (Whnf_ v) = Whnf_ (fmap21 f v) - fmap21 f (Thunk e) = Thunk (f e) - -instance Foldable21 Lazy where - foldMap21 f (Whnf_ v) = foldMap21 f v - foldMap21 f (Thunk e) = f e - -instance Traversable21 Lazy where - traverse21 f (Whnf_ v) = Whnf_ <$> traverse21 f v - traverse21 f (Thunk e) = Thunk <$> f e - - -- | Is the lazy value a variable? getLazyVariable :: (ABT Term abt) => Lazy abt a -> Maybe (Variable a) getLazyVariable e = @@ -481,6 +486,9 @@ data Purity = Pure | Impure | ExpectP deriving (Eq, Read, Show) +-- | A type for tracking the arrays under which the term resides +-- This is used as a binding form when we "lift" transformations +-- (currently only Disintegrate) to work on arrays data Index ast = Ind (Variable 'HNat) (ast 'HNat) instance (ABT Term abt) => Eq (Index (abt '[])) where @@ -489,12 +497,63 @@ instance (ABT Term abt) => Ord (Index (abt '[])) where compare (Ind i _) (Ind j _) = compare i j -- TODO check this -indVar :: Index ast -> Variable 'HNat +indVar :: Index ast -> Variable 'HNat indVar (Ind v _ ) = v - + indSize :: Index ast -> ast 'HNat indSize (Ind _ a) = a +fromIndex :: (ABT Term abt) => Index (abt '[]) -> abt '[] 'HNat +fromIndex (Ind v _) = var v + +-- | Distinguish between variables and heap locations +newtype Location (a :: k) = Location (Variable a) + +instance Show (Sing a) => Show (Location a) where + show (Location v) = show v + +locHint :: Location a -> Text +locHint (Location x) = varHint x + +locType :: Location a -> Sing a +locType (Location x) = varType x + +locEq :: (Show1 (Sing :: k -> *), JmEq1 (Sing :: k -> *)) + => Location (a :: k) + -> Location (b :: k) + -> Maybe (TypeEq a b) +locEq (Location a) (Location b) = varEq a b + +fromLocation :: Location a -> Variable a +fromLocation (Location v) = v + +fromLocations1 :: List1 Location a -> List1 Variable a +fromLocations1 = fmap11 fromLocation + +locations1 :: List1 Variable a -> List1 Location a +locations1 = fmap11 Location + +newtype LAssoc ast = LAssoc (Assoc ast) +newtype LAssocs ast = LAssocs (Assocs ast) + +emptyLAssocs :: LAssocs abt +emptyLAssocs = LAssocs (emptyAssocs) + +singletonLAssocs :: Location a -> f a -> LAssocs f +singletonLAssocs (Location v) e = LAssocs (singletonAssocs v e) + +toLAssocs1 :: List1 Location xs -> List1 ast xs -> LAssocs ast +toLAssocs1 ls es = LAssocs (toAssocs1 (fromLocations1 ls) es) + +insertLAssocs :: LAssocs ast -> LAssocs ast -> LAssocs ast +insertLAssocs (LAssocs a) (LAssocs b) = LAssocs (insertAssocs a b) + +lookupLAssoc :: (Show1 (Sing :: k -> *), JmEq1 (Sing :: k -> *)) + => Location (a :: k) + -> LAssocs ast + -> Maybe (ast a) +lookupLAssoc (Location v) (LAssocs a) = lookupAssoc v a + -- | A single statement in some ambient monad (specified by the @p@ -- type index). In particular, note that the the first argument to -- 'MBind' (or 'Let_') together with the variable bound in the @@ -503,8 +562,11 @@ -- In addition to these binding constructs, we also include a few -- non-binding statements like 'SWeight'. -- +-- Statements are parameterized by the type of the bound element, +-- which (if present) is either a Variable or a Location. +-- -- The semantics of this type are as follows. Let @ss :: [Statement --- abt p]@ be a sequence of statements. We have @Γ@: the collection +-- abt v p]@ be a sequence of statements. We have @Γ@: the collection -- of all free variables that occur in the term expressions in @ss@, -- viewed as a measureable space (namely the product of the measureable -- spaces for each variable). And we have @Δ@: the collection of @@ -512,35 +574,35 @@ -- measurable space. The semantic interpretation of @ss@ is a -- measurable function of type @Γ ':-> M Δ@ where @M@ is either -- @HMeasure@ (if @p ~ 'Impure@) or @Identity@ (if @p ~ 'Pure@). -data Statement :: ([Hakaru] -> Hakaru -> *) -> Purity -> * where +data Statement :: ([Hakaru] -> Hakaru -> *) -> (Hakaru -> *) -> Purity -> * where -- BUG: haddock doesn't like annotations on GADT constructors. So we can't make the constructor descriptions below available to Haddock. -- -- A variable bound by 'MBind' to a measure expression. SBind - :: forall abt (a :: Hakaru) - . {-# UNPACK #-} !(Variable a) + :: forall abt (v :: Hakaru -> *) (a :: Hakaru) + . {-# UNPACK #-} !(v a) -> !(Lazy abt ('HMeasure a)) -> [Index (abt '[])] - -> Statement abt 'Impure + -> Statement abt v 'Impure -- A variable bound by 'Let_' to an expression. SLet - :: forall abt p (a :: Hakaru) - . {-# UNPACK #-} !(Variable a) + :: forall abt p (v :: Hakaru -> *) (a :: Hakaru) + . {-# UNPACK #-} !(v a) -> !(Lazy abt a) -> [Index (abt '[])] - -> Statement abt p + -> Statement abt v p -- A weight; i.e., the first component of each argument to -- 'Superpose_'. This is a statement just so that we can avoid -- needing to atomize the weight itself. SWeight - :: forall abt + :: forall abt (v :: Hakaru -> *) . !(Lazy abt 'HProb) -> [Index (abt '[])] - -> Statement abt 'Impure + -> Statement abt v 'Impure -- A monadic guard statement. If the scrutinee matches the -- pattern, then we bind the variables as usual; otherwise, we @@ -549,12 +611,12 @@ -- /monadic context/. In pure contexts we should be able to -- handle case analysis without putting anything onto the heap. SGuard - :: forall abt (xs :: [Hakaru]) (a :: Hakaru) - . !(List1 Variable xs) + :: forall abt (v :: Hakaru -> *) (xs :: [Hakaru]) (a :: Hakaru) + . !(List1 v xs) -> !(Pattern xs a) -> !(Lazy abt a) -> [Index (abt '[])] - -> Statement abt 'Impure + -> Statement abt v 'Impure -- Some arbitrary pure code. This is a statement just so that we can avoid needing to atomize the stuff in the pure code. -- @@ -562,27 +624,26 @@ -- TODO: generalize to use a 'VarSet' so we can collapse these -- TODO: defunctionalize? These break pretty printing... SStuff0 - :: forall abt + :: forall abt (v :: Hakaru -> *) . (abt '[] 'HProb -> abt '[] 'HProb) -> [Index (abt '[])] - -> Statement abt 'ExpectP + -> Statement abt v 'ExpectP SStuff1 - :: forall abt (a :: Hakaru) - . {-# UNPACK #-} !(Variable a) + :: forall abt (v :: Hakaru -> *) (a :: Hakaru) + . {-# UNPACK #-} !(v a) -> (abt '[] 'HProb -> abt '[] 'HProb) -> [Index (abt '[])] - -> Statement abt 'ExpectP - + -> Statement abt v 'ExpectP -statementVars :: Statement abt p -> VarSet ('KProxy :: KProxy Hakaru) -statementVars (SBind x _ _) = singletonVarSet x -statementVars (SLet x _ _) = singletonVarSet x +statementVars :: Statement abt Location p -> VarSet ('KProxy :: KProxy Hakaru) +statementVars (SBind x _ _) = singletonVarSet (fromLocation x) +statementVars (SLet x _ _) = singletonVarSet (fromLocation x) statementVars (SWeight _ _) = emptyVarSet -statementVars (SGuard xs _ _ _) = toVarSet1 xs +statementVars (SGuard xs _ _ _) = toVarSet1 (fromLocations1 xs) statementVars (SStuff0 _ _) = emptyVarSet -statementVars (SStuff1 x _ _) = singletonVarSet x +statementVars (SStuff1 x _ _) = singletonVarSet (fromLocation x) --- | Is the variable bound by the statement? +-- | Is the Location bound by the statement? -- -- We return @Maybe ()@ rather than @Bool@ because in our primary -- use case we're already in the @Maybe@ monad and so it's easier @@ -590,16 +651,17 @@ -- really rather have the @Bool@, then we can easily change things -- and use some @boolToMaybe@ function to do the coercion wherever -- needed. -isBoundBy :: Variable (a :: Hakaru) -> Statement abt p -> Maybe () -x `isBoundBy` SBind y _ _ = const () <$> varEq x y -x `isBoundBy` SLet y _ _ = const () <$> varEq x y +isBoundBy :: Location (a :: Hakaru) -> Statement abt Location p -> Maybe () +x `isBoundBy` SBind y _ _ = const () <$> locEq x y +x `isBoundBy` SLet y _ _ = const () <$> locEq x y _ `isBoundBy` SWeight _ _ = Nothing x `isBoundBy` SGuard ys _ _ _ = - if memberVarSet x (toVarSet1 ys) -- TODO: just check membership directly, rather than going through VarSet + -- TODO: just check membership directly, rather than going through VarSet + if memberVarSet (fromLocation x) (toVarSet1 (fmap11 fromLocation ys)) then Just () else Nothing _ `isBoundBy` SStuff0 _ _ = Nothing -x `isBoundBy` SStuff1 y _ _ = const () <$> varEq x y +x `isBoundBy` SStuff1 y _ _ = const () <$> locEq x y -- TODO: remove this CPP guard, provided we don't end up with a cyclic dependency... @@ -632,16 +694,16 @@ ppInds :: (ABT Term abt) => [Index (abt '[])] -> PP.Doc ppInds = ppList . map (ppVariable . indVar) -ppStatement :: (ABT Term abt) => Int -> Statement abt p -> PP.Doc +ppStatement :: (ABT Term abt) => Int -> Statement abt Location p -> PP.Doc ppStatement p s = case s of - SBind x e inds -> + SBind (Location x) e inds -> PP.sep $ ppFun p "SBind" [ ppVariable x , PP.sep $ prettyPrec_ 11 e , ppInds inds ] - SLet x e inds -> + SLet (Location x) e inds -> PP.sep $ ppFun p "SLet" [ ppVariable x , PP.sep $ prettyPrec_ 11 e @@ -654,7 +716,7 @@ ] SGuard xs pat e inds -> PP.sep $ ppFun p "SGuard" - [ PP.sep $ ppVariables xs + [ PP.sep $ ppVariables (fromLocations1 xs) , PP.sep $ prettyPrec_ 11 pat , PP.sep $ prettyPrec_ 11 e , ppInds inds @@ -668,7 +730,7 @@ [ PP.text "TODO: ppStatement{SStuff1}" ] -pretty_Statements :: (ABT Term abt) => [Statement abt p] -> PP.Doc +pretty_Statements :: (ABT Term abt) => [Statement abt Location p] -> PP.Doc pretty_Statements [] = PP.text "[]" pretty_Statements (s:ss) = foldl @@ -678,7 +740,7 @@ PP.$+$ PP.text "]" pretty_Statements_withTerm - :: (ABT Term abt) => [Statement abt p] -> abt '[] a -> PP.Doc + :: (ABT Term abt) => [Statement abt Location p] -> abt '[] a -> PP.Doc pretty_Statements_withTerm ss e = pretty_Statements ss PP.$+$ pretty e @@ -693,6 +755,29 @@ #endif + +----------------------------------------------------------------- +-- | A function for evaluating any term to weak-head normal form. +type TermEvaluator abt m = + forall a. abt '[] a -> m (Whnf abt a) + +-- | A function for \"performing\" an 'HMeasure' monadic action. +-- This could mean actual random sampling, or simulated sampling +-- by generating a new term and returning the newly bound variable, +-- or anything else. +type MeasureEvaluator abt m = + forall a. abt '[] ('HMeasure a) -> m (Whnf abt a) + +-- | A function for evaluating any case-expression to weak-head +-- normal form. +type CaseEvaluator abt m = + forall a b. abt '[] a -> [Branch a abt b] -> m (Whnf abt b) + +-- | A function for evaluating any variable to weak-head normal form. +type VariableEvaluator abt m = + forall a. Variable a -> m (Whnf abt a) + + ---------------------------------------------------------------- -- | This class captures the monadic operations needed by the -- 'evaluate' function in "Language.Hakaru.Lazy". @@ -710,25 +795,26 @@ -- | Internal function for renaming the variables bound by a -- statement. We return the renamed statement along with a substitution -- for mapping the old variable names to their new variable names. - freshenStatement - :: Statement abt p - -> m (Statement abt p, Assocs (Variable :: Hakaru -> *)) - freshenStatement s = + freshLocStatement + :: Statement abt Variable p + -> m (Statement abt Location p, Assocs (Variable :: Hakaru -> *)) + freshLocStatement s = case s of - SWeight _ _ -> return (s, mempty) + SWeight w e -> return (SWeight w e, mempty) SBind x body i -> do x' <- freshenVar x - return (SBind x' body i, singletonAssocs x x') - SLet x body i -> do + return (SBind (Location x') body i, singletonAssocs x x') + SLet x body i -> do x' <- freshenVar x - return (SLet x' body i, singletonAssocs x x') + return (SLet (Location x') body i, singletonAssocs x x') SGuard xs pat scrutinee i -> do xs' <- freshenVars xs - return (SGuard xs' pat scrutinee i, toAssocs1 xs xs') - SStuff0 _ _ -> return (s, mempty) + return (SGuard (locations1 xs') pat scrutinee i, + toAssocs1 xs xs') + SStuff0 e e' -> return (SStuff0 e e', mempty) SStuff1 x f i -> do x' <- freshenVar x - return (SStuff1 x' f i, singletonAssocs x x') + return (SStuff1 (Location x') f i, singletonAssocs x x') -- | Returns the current Indices. Currently, this is only @@ -742,7 +828,7 @@ -- because it may allow confusion between variables with the -- same name but different scopes (thus, may allow variable -- capture). Prefer using 'push_', 'push', or 'pushes'. - unsafePush :: Statement abt p -> m () + unsafePush :: Statement abt Location p -> m () -- | Call 'unsafePush' repeatedly. Is part of the class since -- we may be able to do this more efficiently than actually @@ -750,7 +836,7 @@ -- -- N.B., this should push things in the same order as 'pushes' -- does. - unsafePushes :: [Statement abt p] -> m () + unsafePushes :: [Statement abt Location p] -> m () unsafePushes = mapM_ unsafePush -- | Look for the statement @s@ binding the variable. If found, @@ -782,12 +868,138 @@ -- 'Statement'. Perhaps we should have an alternative statement -- type which exposes the existential? select - :: Variable (a :: Hakaru) - -> (Statement abt p -> Maybe (m r)) + :: Location (a :: Hakaru) + -> (Statement abt Location p -> Maybe (m r)) -> m (Maybe r) + substVar :: Variable a -> abt '[] a + -> (forall b'. Variable b' -> m (abt '[] b')) + substVar x e = return . var + extFreeVars :: abt xs a -> m (VarSet (KindOf a)) + extFreeVars e = return (freeVars e) + + + -- The first argument to @evaluateCase@ will be the + -- 'TermEvaluator' we're constructing (thus tying the knot). + evaluateCase :: TermEvaluator abt m -> CaseEvaluator abt m + {-# INLINE evaluateCase #-} + evaluateCase = defaultCaseEvaluator + + + -- TODO: figure out how to abstract this so it can be reused by + -- 'constrainValue'. Especially the 'SBranch case of 'step' + -- TODO: we could speed up the case for free variables by having + -- the 'Context' also keep track of the largest free var. That way, + -- we can just check up front whether @varID x < nextFreeVarID@. + -- Of course, we'd have to make sure we've sufficiently renamed all + -- bound variables to be above @nextFreeVarID@; but then we have to + -- do that anyways. + evaluateVar :: MeasureEvaluator abt m + -> TermEvaluator abt m + -> VariableEvaluator abt m + evaluateVar perform evaluate_ = \x -> + -- If we get 'Nothing', then it turns out @x@ is a free variable + fmap (maybe (Neutral $ var x) id) . select (Location x) $ \s -> + case s of + SBind y e i -> do + Refl <- locEq (Location x) y + Just $ do + w <- perform $ caseLazy e fromWhnf id + unsafePush (SLet (Location x) (Whnf_ w) i) +#ifdef __TRACE_DISINTEGRATE__ + trace ("-- updated " + ++ show (ppStatement 11 s) + ++ " to " + ++ show (ppStatement 11 (SLet (Location x) (Whnf_ w) i)) + ) $ return () +#endif + return w + SLet y e i -> do + Refl <- locEq (Location x) y + Just $ do + w <- caseLazy e return evaluate_ + unsafePush (SLet (Location x) (Whnf_ w) i) + return w + -- These two don't bind any variables, so they definitely + -- can't match. + SWeight _ _ -> Nothing + SStuff0 _ _ -> Nothing + -- These two do bind variables, but there's no expression we + -- can return for them because the variables are + -- untouchable\/abstract. + SStuff1 _ _ _ -> Just . return . Neutral $ var x + SGuard ys pat scrutinee i -> Just . return . Neutral $ var x + + +-- | A simple 'CaseEvaluator' which uses the 'DatumEvaluator' to +-- force the scrutinee, and if 'matchBranches' succeeds then we +-- call the 'TermEvaluator' to continue evaluating the body of the +-- matched branch. If we 'GotStuck' then we return a 'Neutral' term +-- of the case expression itself (n.b, any side effects from having +-- called the 'DatumEvaluator' will still persist when returning +-- this neutral term). If we didn't get stuck and yet none of the +-- branches matches, then we throw an exception. +defaultCaseEvaluator + :: forall abt m p + . (ABT Term abt, EvaluationMonad abt m p) + => TermEvaluator abt m + -> CaseEvaluator abt m +{-# INLINE defaultCaseEvaluator #-} +defaultCaseEvaluator evaluate_ = evaluateCase_ + where + -- TODO: At present, whenever we residualize a case expression we'll + -- generate a 'Neutral' term which will, when run, repeat the work + -- we're doing in the evaluation here. We could eliminate this + -- redundancy by introducing a new variable for @e@ each time this + -- function is called--- if only we had some way of getting those + -- variables put into the right place for when we residualize the + -- original scrutinee... + -- + -- N.B., 'DatumEvaluator' is a rank-2 type so it requires a signature + evaluateDatum :: DatumEvaluator (abt '[]) m + evaluateDatum e = viewWhnfDatum <$> evaluate_ e + + evaluateCase_ :: CaseEvaluator abt m + evaluateCase_ e bs = do + match <- matchBranches evaluateDatum e bs + case match of + Nothing -> + -- TODO: print more info about where this error + -- happened + -- + -- TODO: rather than throwing a Haskell error, + -- instead capture the possibility of failure in + -- the 'EvaluationMonad' monad. + error "defaultCaseEvaluator: non-exhaustive patterns in case!" + Just GotStuck -> + return . Neutral . syn $ Case_ e bs + Just (Matched ss body) -> + pushes (toVarStatements ss) body >>= evaluate_ + + +toVarStatements :: Assocs (abt '[]) -> [Statement abt Variable p] +toVarStatements = map (\(Assoc x e) -> SLet x (Thunk e) []) . + fromAssocs + +extSubst + :: forall abt a xs b m p. (EvaluationMonad abt m p) + => Variable a + -> abt '[] a + -> abt xs b + -> m (abt xs b) +extSubst x e = substM x e (substVar x e) + +extSubsts + :: forall abt a xs m p. (EvaluationMonad abt m p) + => Assocs (abt '[]) + -> abt xs a + -> m (abt xs a) +extSubsts rho0 e0 = + F.foldlM (\e (Assoc x v) -> extSubst x v e) e0 (unAssocs rho0) + + -- TODO: define a new NameSupply monad in "Language.Hakaru.Syntax.Variable" for encapsulating these four fresh(en) functions? @@ -858,6 +1070,21 @@ return $ Ind x s +-- | Given a location, return a new Location with the same hint +-- and type but with a fresh ID +freshenLoc :: (EvaluationMonad abt m p) + => Location (a :: Hakaru) -> m (Location a) +freshenLoc (Location x) = Location <$> freshenVar x + +-- | Call `freshenLoc` repeatedly +freshenLocs :: (EvaluationMonad abt m p) + => List1 Location (ls :: [Hakaru]) + -> m (List1 Location ls) +freshenLocs Nil1 = return Nil1 +freshenLocs (Cons1 l ls) = Cons1 <$> freshenLoc l <*> freshenLocs ls + + + -- | Add a statement to the top of the context, renaming any variables -- the statement binds and returning the substitution mapping the -- old variables to the new ones. This is safer than 'unsafePush' @@ -867,10 +1094,10 @@ -- 'pushes' instead. push_ :: (ABT Term abt, EvaluationMonad abt m p) - => Statement abt p + => Statement abt Variable p -> m (Assocs (Variable :: Hakaru -> *)) push_ s = do - (s',rho) <- freshenStatement s + (s',rho) <- freshLocStatement s unsafePush s' return rho @@ -887,13 +1114,13 @@ -- than applying it to the term before calling the continuation. push :: (ABT Term abt, EvaluationMonad abt m p) - => Statement abt p -- ^ the statement to push + => Statement abt Variable p -- ^ the statement to push -> abt xs a -- ^ the \"rest\" of the term - -> (abt xs a -> m r) -- ^ what to do with the renamed \"rest\" - -> m r -- ^ the final result -push s e k = do + -- -> (abt xs a -> m r) -- ^ what to do with the renamed \"rest\" + -> m (abt xs a) -- ^ the final result +push s e = do rho <- push_ s - k (renames rho e) + return (renames rho e) -- | Call 'push' repeatedly. (N.B., is more efficient than actually @@ -902,14 +1129,14 @@ -- pushed last and is the closest in the final context. pushes :: (ABT Term abt, EvaluationMonad abt m p) - => [Statement abt p] -- ^ the statements to push + => [Statement abt Variable p] -- ^ the statements to push -> abt xs a -- ^ the \"rest\" of the term - -> (abt xs a -> m r) -- ^ what to do with the renamed \"rest\" - -> m r -- ^ the final result -pushes ss e k = do + -- -> (abt xs a -> m r) -- ^ what to do with the renamed \"rest\" + -> m (abt xs a) -- ^ the final result +pushes ss e = do -- TODO: is 'foldlM' the right one? or do we want 'foldrM'? rho <- F.foldlM (\rho s -> mappend rho <$> push_ s) mempty ss - k (renames rho e) + return (renames rho e) ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/haskell/Language/Hakaru/Expect.hs b/haskell/Language/Hakaru/Expect.hs --- a/haskell/Language/Hakaru/Expect.hs +++ b/haskell/Language/Hakaru/Expect.hs @@ -98,7 +98,7 @@ residualizeExpect e = do -- BUG: is this what we really mean? or do we actually mean the old 'emit' version? x <- freshVar Text.empty (sUnMeasure $ typeOf e) - unsafePush (SStuff1 x (\c -> + unsafePush (SStuff1 (Location x) (\c -> syn (AST.Expect :$ e :* bind x c :* End)) []) return $ var x {- @@ -165,7 +165,7 @@ -- TODO: do we want to move this to the public API of -- "Language.Hakaru.Evaluation.DisintegrationMonad"? #ifdef __TRACE_DISINTEGRATE__ -getStatements :: Expect abt [Statement abt 'ExpectP] +getStatements :: Expect abt [Statement abt Location 'ExpectP] getStatements = Expect $ \c h -> c (statements h) h #endif @@ -189,9 +189,9 @@ case t of Case_ e1 bs -> expectCase e1 bs _ -> residualizeExpect e' - Head_ (WLiteral _) -> error "expect: the impossible happened" - Head_ (WCoerceTo _ _) -> error "expect: the impossible happened" - Head_ (WUnsafeFrom _ _) -> error "expect: the impossible happened" + Head_ (WLiteral _) -> error "expect: the impossible happened" + Head_ (WCoerceTo _ _) -> error "expect: the impossible happened" + Head_ (WUnsafeFrom _ _) -> error "expect: the impossible happened" Head_ (WMeasureOp o es) -> expectMeasureOp o es Head_ (WDirac e1) -> return e1 Head_ (WMBind e1 e2) -> do @@ -229,16 +229,16 @@ ss <- Expect $ \c h -> c (statements h) (h {statements = []}) F.traverse_ step (reverse ss) -- TODO: use composition tricks to avoid reversing @ss@ where - step :: Statement abt 'ExpectP -> Expect abt () + step :: Statement abt Location 'ExpectP -> Expect abt () step s = #ifdef __TRACE_DISINTEGRATE__ trace ("\n-- emitExpectListContext: " ++ show (ppStatement 0 s)) $ #endif case s of - SLet x body _ -> + SLet l body _ -> -- TODO: be smart about dropping unused let-bindings and inlining trivial let-bindings Expect $ \c h -> - syn (Let_ :$ fromLazy body :* bind x (c () h) :* End) + syn (Let_ :$ fromLazy body :* bind (fromLocation l) (c () h) :* End) SStuff0 f _ -> Expect $ \c h -> f (c () h) SStuff1 _ f _ -> Expect $ \c h -> f (c () h) @@ -250,7 +250,7 @@ -> Expect abt (Variable 'HReal) pushIntegrate lo hi = do x <- freshVar Text.empty SReal - unsafePush (SStuff1 x (\c -> + unsafePush (SStuff1 (Location x) (\c -> syn (Integrate :$ lo :* hi :* bind x c :* End)) []) return x {- @@ -268,7 +268,7 @@ -> Expect abt (Variable a) pushSummate lo hi = do x <- freshVar Text.empty sing - unsafePush (SStuff1 x (\c -> + unsafePush (SStuff1 (Location x) (\c -> syn (Summate hDiscrete hSemiring :$ lo :* hi :* bind x c :* End)) []) return x @@ -285,7 +285,7 @@ pushLet e = caseVarSyn e return $ \_ -> do x <- freshVar Text.empty (typeOf e) - unsafePush (SStuff1 x (\c -> + unsafePush (SStuff1 (Location x) (\c -> syn (Let_ :$ e :* bind x c :* End)) []) return x {- diff --git a/haskell/Language/Hakaru/Maple.hs b/haskell/Language/Hakaru/Maple.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Maple.hs @@ -0,0 +1,146 @@ +{-# LANGUAGE FlexibleInstances + , FlexibleContexts + , DeriveDataTypeable + , DataKinds + , ScopedTypeVariables + , RecordWildCards + , ViewPatterns + , LambdaCase + , DeriveFunctor, DeriveFoldable, DeriveTraversable + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- | +-- Module : Language.Hakaru.Maple +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Stability : experimental +-- Portability : GHC-only +-- +-- Take strings from Maple and interpret them in Haskell (Hakaru), +-- in a type-safe way. +---------------------------------------------------------------- +module Language.Hakaru.Maple + ( MapleException(..) + , MapleOptions(..) + , defaultMapleOptions + , sendToMaple, sendToMaple' + , maple + ) where + +import Control.Exception +import Control.Monad (when) + +import qualified Language.Hakaru.Pretty.Maple as Maple + +import Language.Hakaru.Parser.Maple +import Language.Hakaru.Parser.AST (Name) +import qualified Language.Hakaru.Parser.SymbolResolve as SR (resolveAST', fromVarSet) + +import Language.Hakaru.Types.Sing +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.TypeCheck +import Language.Hakaru.Syntax.TypeOf +import Language.Hakaru.Syntax.Command + +import Language.Hakaru.Evaluation.ConstantPropagation + +import Data.Typeable (Typeable) + +import System.MapleSSH (maple) +import System.IO +import Data.Text (pack) +import qualified Data.Map as M +import Data.List (intercalate) + +import Data.Foldable (Foldable) +import Data.Traversable (Traversable) + +---------------------------------------------------------------- +data MapleException + = MapleInterpreterException String String + | MapleInputTypeMismatch String String + | MapleUnknownCommand String + deriving Typeable + +instance Exception MapleException + +-- Maple prints errors with "cursors" (^) which point to the specific position +-- of the error on the line above. The derived show instance doesn't preserve +-- positioning of the cursor. +instance Show MapleException where + show (MapleInterpreterException toMaple_ fromMaple) = + "MapleException:\n" ++ fromMaple ++ + "\nafter sending to Maple:\n" ++ toMaple_ + show (MapleInputTypeMismatch command ty) = + concat["Maple command ", command, " does not take input of type ", ty] + show (MapleUnknownCommand command) = + concat["Maple command ", command, " does not exist"] + +data MapleOptions nm = MapleOptions + { command :: nm + , debug :: Bool + , timelimit :: Int + , extraOpts :: M.Map String String + } deriving (Functor, Foldable, Traversable) + +defaultMapleOptions :: MapleOptions () +defaultMapleOptions = MapleOptions + { command = () + , debug = False + , timelimit = 90 + , extraOpts = M.empty } + +sendToMaple' + :: ABT Term (abt Term) + => MapleOptions String + -> TypedAST (abt Term) + -> IO (TypedAST (abt Term)) +sendToMaple' MapleOptions{..} (TypedAST ty expr) = + commandFromName command ty $ \case + Left True -> throw $ MapleInputTypeMismatch command (show ty) + Left False -> throw $ MapleUnknownCommand command + Right (c, ty_o) -> fmap (TypedAST ty_o) (sendToMaple MapleOptions{command=c,..} expr) + +sendToMaple + :: (ABT Term abt) + => MapleOptions (CommandType c i o) + -> abt '[] i + -> IO (abt '[] o) +sendToMaple MapleOptions{..} e = do + let typ_in = typeOf e + typ_out = commandIsType command typ_in + optStr (k,v) = concat["_",k,"=",v] + optsStr = + intercalate "," $ + map optStr $ M.assocs $ + M.insert "command" (ssymbolVal(nameOfCommand command)) extraOpts + toMaple_ = "use Hakaru, NewSLO in timelimit(" + ++ show timelimit ++ ", RoundTrip(" + ++ Maple.pretty e ++ ", " ++ Maple.mapleType typ_in (", " + ++ optsStr ++ ")) end use;") + when debug (hPutStrLn stderr ("Sent to Maple:\n" ++ toMaple_)) + fromMaple <- maple toMaple_ + case fromMaple of + '_':'I':'n':'e':'r':'t':_ -> do + when debug $ do + ret <- maple ("FromInert(" ++ fromMaple ++ ")") + hPutStrLn stderr ("Returning from Maple:\n" ++ ret) + either (throw . MapleInterpreterException toMaple_) + (return . constantPropagation) $ do + past <- leftShow $ parseMaple (pack fromMaple) + let m = checkType typ_out + (SR.resolveAST' (getNames e) (maple2AST past)) + leftShow $ unTCM m (freeVars e) Nothing UnsafeMode + _ -> throw (MapleInterpreterException toMaple_ fromMaple) + +leftShow :: forall b c. Show b => Either b c -> Either String c +leftShow (Left err) = Left (show err) +leftShow (Right x) = Right x + +getNames :: ABT Term abt => abt '[] a -> [Name] +getNames = SR.fromVarSet . freeVars + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/haskell/Language/Hakaru/Observe.hs b/haskell/Language/Hakaru/Observe.hs --- a/haskell/Language/Hakaru/Observe.hs +++ b/haskell/Language/Hakaru/Observe.hs @@ -79,6 +79,7 @@ -- This function can't inspect a variable due to -- calls to subst that happens in Let_ and Bind_ +observeVar :: Variable a -> r observeVar = error "observe can only be applied measure primitives" observeMeasureOp diff --git a/haskell/Language/Hakaru/Parser/AST.hs b/haskell/Language/Hakaru/Parser/AST.hs --- a/haskell/Language/Hakaru/Parser/AST.hs +++ b/haskell/Language/Hakaru/Parser/AST.hs @@ -5,6 +5,7 @@ , ExistentialQuantification , StandaloneDeriving , TypeFamilies + , TypeOperators , OverloadedStrings #-} @@ -98,7 +99,7 @@ data NaryOp = And | Or | Xor - | Iff | Min | Max + | Iff | Min | Max | Sum | Prod deriving (Eq, Show) @@ -110,9 +111,17 @@ | TypeFun TypeAST' TypeAST' deriving (Eq, Show) +data Reducer' a + = R_Fanout (Reducer' a) (Reducer' a) + | R_Index a (AST' a) (AST' a) (Reducer' a) + | R_Split (AST' a) (Reducer' a) (Reducer' a) + | R_Nop + | R_Add (AST' a) + deriving (Eq, Show) + data AST' a = Var a - | Lam a TypeAST' (AST' a) + | Lam a TypeAST' (AST' a) | App (AST' a) (AST' a) | Let a (AST' a) (AST' a) | If (AST' a) (AST' a) (AST' a) @@ -124,6 +133,7 @@ | Empty | Pair (AST' a) (AST' a) | Array a (AST' a) (AST' a) + | ArrayLiteral [AST' a] | Index (AST' a) (AST' a) | Case (AST' a) [(Branch' a)] -- match | Dirac (AST' a) @@ -133,10 +143,11 @@ | Integrate a (AST' a) (AST' a) (AST' a) | Summate a (AST' a) (AST' a) (AST' a) | Product a (AST' a) (AST' a) (AST' a) + | Bucket a (AST' a) (AST' a) (Reducer' a) | Expect a (AST' a) (AST' a) | Observe (AST' a) (AST' a) | Msum [AST' a] - | Data a [TypeAST'] + | Data a [a] [TypeAST'] (AST' a) | WithMeta (AST' a) SourceSpan deriving (Show) @@ -166,6 +177,7 @@ (Array e1 e2 e3) == (Array e1' e2' e3') = e1 == e1' && e2 == e2' && e3 == e3' + (ArrayLiteral es) == (ArrayLiteral es') = es == es' (Index e1 e2) == (Index e1' e2') = e1 == e1' && e2 == e2' (Case e1 bs) == (Case e1' bs') = e1 == e1' && @@ -193,18 +205,24 @@ b == b' && c == c' && d == d' + (Bucket a b c d) == (Bucket a' b' c' d') = a == a' && + b == b' && + c == c' && + d == d' (Expect e1 e2 e3) == (Expect e1' e2' e3') = e1 == e1' && e2 == e2' && e3 == e3' (Observe e1 e2) == (Observe e1' e2') = e1 == e1' && e2 == e2' (Msum es) == (Msum es') = es == es' - (Data e1 ts) == (Data e1' ts') = e1 == e1' && - ts == ts' + (Data n ft ts e) == (Data n' ft' ts' e') = n == n' && + ft == ft' && + ts == ts' && + e == e' (WithMeta e1 _ ) == e2 = e1 == e2 e1 == (WithMeta e2 _) = e1 == e2 _ == _ = False - + data Import a = Import a deriving (Eq, Show) data ASTWithImport' a = ASTWithImport' [Import a] (AST' a) @@ -318,6 +336,41 @@ fdfun f' (Ident e) = f' e +data Reducer (xs :: [Untyped]) + (abt :: [Untyped] -> Untyped -> *) + (a :: Untyped) where + R_Fanout_ :: Reducer xs abt 'U + -> Reducer xs abt 'U + -> Reducer xs abt 'U + R_Index_ :: Variable 'U -- HACK: Shouldn't need to pass this argument + -> abt xs 'U + -> abt ( 'U ': xs) 'U + -> Reducer ( 'U ': xs) abt 'U + -> Reducer xs abt 'U + R_Split_ :: abt ( 'U ': xs) 'U + -> Reducer xs abt 'U + -> Reducer xs abt 'U + -> Reducer xs abt 'U + R_Nop_ :: Reducer xs abt 'U + R_Add_ :: abt ( 'U ': xs) 'U + -> Reducer xs abt 'U + +instance Functor21 (Reducer xs) where + fmap21 f (R_Fanout_ r1 r2) = R_Fanout_ (fmap21 f r1) (fmap21 f r2) + fmap21 f (R_Index_ bv e1 e2 r1) = R_Index_ bv (f e1) (f e2) (fmap21 f r1) + fmap21 f (R_Split_ e1 r1 r2) = R_Split_ (f e1) (fmap21 f r1) (fmap21 f r2) + fmap21 _ R_Nop_ = R_Nop_ + fmap21 f (R_Add_ e1) = R_Add_ (f e1) + +instance Foldable21 (Reducer xs) where + foldMap21 f (R_Fanout_ r1 r2) = foldMap21 f r1 `mappend` foldMap21 f r2 + foldMap21 f (R_Index_ _ e1 e2 r1) = + f e1 `mappend` f e2 `mappend` foldMap21 f r1 + foldMap21 f (R_Split_ e1 r1 r2) = + f e1 `mappend` foldMap21 f r1 `mappend` foldMap21 f r2 + foldMap21 _ R_Nop_ = mempty + foldMap21 f (R_Add_ e1) = f e1 + -- | The kind containing exactly one type. data Untyped = U deriving (Read, Show) @@ -348,33 +401,35 @@ nameToVar (Name i h) = Variable h i SU data Term :: ([Untyped] -> Untyped -> *) -> Untyped -> * where - Lam_ :: SSing -> abt '[ 'U ] 'U -> Term abt 'U - App_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U - Let_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Ann_ :: SSing -> abt '[] 'U -> Term abt 'U - CoerceTo_ :: Some2 Coercion -> abt '[] 'U -> Term abt 'U - UnsafeTo_ :: Some2 Coercion -> abt '[] 'U -> Term abt 'U - PrimOp_ :: PrimOp -> [abt '[] 'U] -> Term abt 'U - ArrayOp_ :: ArrayOp -> [abt '[] 'U] -> Term abt 'U - MeasureOp_ :: SomeOp MeasureOp -> [abt '[] 'U] -> Term abt 'U - NaryOp_ :: NaryOp -> [abt '[] 'U] -> Term abt 'U - Literal_ :: Some1 Literal -> Term abt 'U - Empty_ :: Term abt 'U - Pair_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U - Array_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Datum_ :: Datum abt -> Term abt 'U - Case_ :: abt '[] 'U -> [Branch_ abt] -> Term abt 'U - Dirac_ :: abt '[] 'U -> Term abt 'U - MBind_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Plate_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Chain_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Integrate_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Summate_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Product_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Expect_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U - Observe_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U - Superpose_ :: L.NonEmpty (abt '[] 'U, abt '[] 'U) -> Term abt 'U - Reject_ :: Term abt 'U + Lam_ :: SSing -> abt '[ 'U ] 'U -> Term abt 'U + App_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U + Let_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Ann_ :: SSing -> abt '[] 'U -> Term abt 'U + CoerceTo_ :: Some2 Coercion -> abt '[] 'U -> Term abt 'U + UnsafeTo_ :: Some2 Coercion -> abt '[] 'U -> Term abt 'U + PrimOp_ :: PrimOp -> [abt '[] 'U] -> Term abt 'U + ArrayOp_ :: ArrayOp -> [abt '[] 'U] -> Term abt 'U + MeasureOp_ :: SomeOp MeasureOp -> [abt '[] 'U] -> Term abt 'U + NaryOp_ :: NaryOp -> [abt '[] 'U] -> Term abt 'U + Literal_ :: Some1 Literal -> Term abt 'U + Empty_ :: Term abt 'U + Pair_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U + Array_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + ArrayLiteral_ :: [abt '[] 'U] -> Term abt 'U + Datum_ :: Datum abt -> Term abt 'U + Case_ :: abt '[] 'U -> [Branch_ abt] -> Term abt 'U + Dirac_ :: abt '[] 'U -> Term abt 'U + MBind_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Plate_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Chain_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Integrate_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Summate_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Product_ :: abt '[] 'U -> abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Bucket_ :: abt '[] 'U -> abt '[] 'U -> Reducer xs abt 'U -> Term abt 'U + Expect_ :: abt '[] 'U -> abt '[ 'U ] 'U -> Term abt 'U + Observe_ :: abt '[] 'U -> abt '[] 'U -> Term abt 'U + Superpose_ :: L.NonEmpty (abt '[] 'U, abt '[] 'U) -> Term abt 'U + Reject_ :: Term abt 'U -- TODO: instance of Traversable21 for Term @@ -393,6 +448,7 @@ fmap21 _ Empty_ = Empty_ fmap21 f (Pair_ e1 e2) = Pair_ (f e1) (f e2) fmap21 f (Array_ e1 e2) = Array_ (f e1) (f e2) + fmap21 f (ArrayLiteral_ es) = ArrayLiteral_ (fmap f es) fmap21 f (Datum_ d) = Datum_ (fmapDatum f d) fmap21 f (Case_ e1 bs) = Case_ (f e1) (fmap (fmapBranch f) bs) fmap21 f (Dirac_ e1) = Dirac_ (f e1) @@ -402,6 +458,7 @@ fmap21 f (Integrate_ e1 e2 e3) = Integrate_ (f e1) (f e2) (f e3) fmap21 f (Summate_ e1 e2 e3) = Summate_ (f e1) (f e2) (f e3) fmap21 f (Product_ e1 e2 e3) = Product_ (f e1) (f e2) (f e3) + fmap21 f (Bucket_ e1 e2 e3) = Bucket_ (f e1) (f e2) (fmap21 f e3) fmap21 f (Expect_ e1 e2) = Expect_ (f e1) (f e2) fmap21 f (Observe_ e1 e2) = Observe_ (f e1) (f e2) fmap21 f (Superpose_ es) = Superpose_ (L.map (f *** f) es) @@ -422,6 +479,7 @@ foldMap21 _ Empty_ = mempty foldMap21 f (Pair_ e1 e2) = f e1 `mappend` f e2 foldMap21 f (Array_ e1 e2) = f e1 `mappend` f e2 + foldMap21 f (ArrayLiteral_ es) = F.foldMap f es foldMap21 f (Datum_ d) = foldDatum f d foldMap21 f (Case_ e1 bs) = f e1 `mappend` F.foldMap (foldBranch f) bs foldMap21 f (Dirac_ e1) = f e1 @@ -431,18 +489,20 @@ foldMap21 f (Integrate_ e1 e2 e3) = f e1 `mappend` f e2 `mappend` f e3 foldMap21 f (Summate_ e1 e2 e3) = f e1 `mappend` f e2 `mappend` f e3 foldMap21 f (Product_ e1 e2 e3) = f e1 `mappend` f e2 `mappend` f e3 + foldMap21 f (Bucket_ e1 e2 e3) = f e1 `mappend` f e2 `mappend` foldMap21 f e3 foldMap21 f (Expect_ e1 e2) = f e1 `mappend` f e2 foldMap21 f (Observe_ e1 e2) = f e1 `mappend` f e2 foldMap21 f (Superpose_ es) = F.foldMap (\(e1,e2) -> f e1 `mappend` f e2) es foldMap21 _ Reject_ = mempty -type AST = MetaABT SourceSpan Term '[] 'U -type MetaTerm = Term (MetaABT SourceSpan Term) 'U -type Branch = Branch_ (MetaABT SourceSpan Term) +type U_ABT = MetaABT SourceSpan Term +type AST = U_ABT '[] 'U +type MetaTerm = Term U_ABT 'U +type Branch = Branch_ U_ABT -type DFun_ = DFun (MetaABT SourceSpan Term) -type DStruct_ = DStruct (MetaABT SourceSpan Term) -type DCode_ = DCode (MetaABT SourceSpan Term) +type DFun_ = DFun U_ABT +type DStruct_ = DStruct U_ABT +type DCode_ = DCode U_ABT ---------------------------------------------------------------- ---------------------------------------------------------- fin. diff --git a/haskell/Language/Hakaru/Parser/Import.hs b/haskell/Language/Hakaru/Parser/Import.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Parser/Import.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE CPP, OverloadedStrings #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Language.Hakaru.Parser.Import where + +import Language.Hakaru.Parser.AST +import Language.Hakaru.Parser.Parser + +import Control.Monad.Trans.Except +import Control.Monad.IO.Class +import qualified Data.Text as T +import qualified Data.Text.IO as IO +import Text.Parsec hiding (Empty) + +replaceBody :: AST' T.Text -> AST' T.Text -> AST' T.Text +replaceBody e1 e2 = + case e1 of + Let x e3 e4 -> Let x e3 (replaceBody e4 e2) + Ann e3 t -> Ann (replaceBody e3 e2) t + WithMeta e3 s -> WithMeta (replaceBody e3 e2) s + _ -> e2 + +expandImports + :: ASTWithImport' T.Text + -> ExceptT ParseError IO (AST' T.Text) +expandImports (ASTWithImport' (Import i:is) ast) = do + file <- liftIO . IO.readFile . T.unpack $ T.append i ".hk" + astIm <- ExceptT . return $ parseHakaruWithImports file + ast' <- expandImports astIm + expandImports (ASTWithImport' is (replaceBody ast' ast)) +expandImports (ASTWithImport' [] ast) = return ast diff --git a/haskell/Language/Hakaru/Parser/Maple.hs b/haskell/Language/Hakaru/Parser/Maple.hs --- a/haskell/Language/Hakaru/Parser/Maple.hs +++ b/haskell/Language/Hakaru/Parser/Maple.hs @@ -130,6 +130,15 @@ InertName <$> (text "_Inert_NAME" *> parens stringLiteral) +localname :: Parser InertExpr +localname = + InertName + <$> (text "_Inert_LOCALNAME" + *> parens + ( stringLiteral + <* comma + <* integer)) + assignedname :: Parser InertExpr assignedname = InertName @@ -259,6 +268,7 @@ <|> try noteq <|> try assignedname <|> try assignedlocalname + <|> try localname <|> try expseq <|> try intpos <|> try intneg @@ -296,6 +306,10 @@ ULiteral . Real $ fromInteger a * (10 ^ b) maple2AST (InertArgs Func + [InertName "Let", InertArgs ExpSeq [e1, InertName x, e2]]) = + Let x (maple2AST e1) (maple2AST e2) + +maple2AST (InertArgs Func [InertName "Bind", InertArgs ExpSeq [e1, InertName x, e2]]) = Bind x (maple2AST e1) (maple2AST e2) @@ -348,9 +362,18 @@ [InertName "piecewise", InertArgs ExpSeq es]) = go es where go [e1,e2] = If (maple2AST e1) (maple2AST e2) (ULiteral (Nat 0)) go [e1,e2,e3] = If (maple2AST e1) (maple2AST e2) (maple2AST e3) + -- BUG! piecewise(a Reducer' Text +maple2ReducerAST + (InertArgs Func + [ InertName "Fanout" + , InertArgs ExpSeq [ e1, e2 ]]) = + R_Fanout (maple2ReducerAST e1) (maple2ReducerAST e2) + +maple2ReducerAST + (InertArgs Func + [ InertName "Index" + , InertArgs ExpSeq [ e1, InertName x, e2, e3]]) = + R_Index x (maple2AST e1) (maple2AST e2) (maple2ReducerAST e3) + +maple2ReducerAST + (InertArgs Func + [ InertName "Split" + , InertArgs ExpSeq [ e1, e2, e3]]) = + R_Split (maple2AST e1) (maple2ReducerAST e2) (maple2ReducerAST e3) + +maple2ReducerAST + (InertArgs Func + [ InertName "Nop" + , InertArgs ExpSeq []]) = R_Nop + +maple2ReducerAST + (InertArgs Func + [ InertName "Add" + , InertArgs ExpSeq [e1]]) = R_Add (maple2AST e1) + mapleDatum2AST :: Text -> InertExpr -> AST' Text mapleDatum2AST h d = case (h, maple2DCode d) of ("pair", [x,y]) -> Pair x y @@ -485,16 +566,18 @@ InertArgs ExpSeq []]) = TypeVar "int" maple2Type (InertArgs Func - [InertName "HReal", + [InertName nm, InertArgs ExpSeq [InertArgs Func [InertName "Bound", InertArgs ExpSeq [InertName ">=",InertNum Pos 0]]]]) + | nm `elem` [ "HReal", "AlmostEveryReal" ] = TypeVar "prob" maple2Type (InertArgs Func - [InertName "HReal", + [InertName nm, InertArgs ExpSeq []]) + | nm `elem` [ "HReal", "AlmostEveryReal" ] = TypeVar "real" maple2Type (InertArgs Func diff --git a/haskell/Language/Hakaru/Parser/Parser.hs b/haskell/Language/Hakaru/Parser/Parser.hs --- a/haskell/Language/Hakaru/Parser/Parser.hs +++ b/haskell/Language/Hakaru/Parser/Parser.hs @@ -26,7 +26,8 @@ ops, types, names :: [String] -ops = ["+","*","-","^", "**", ":",".", "<~","==", "=", "_", "<|>"] +ops = ["+","*","-","^", "**", ":",".", "<~", + "==", "=", "_", "<|>", "&&", "||"] types = ["->"] names = ["def", "fn", "if", "else", "∞", "expect", "observe", "return", "match", "integrate", "summate", "product", @@ -74,7 +75,7 @@ float :: Parser Rational float = (decimal >>= fractExponent) <* whiteSpace - + fractFloat :: Integer -> Parser (Either Integer Rational) fractFloat n = fractExponent n >>= return . Right @@ -176,13 +177,14 @@ | s == ">=" = NaryOp Or [ app2 "less" y x , app2 "equal" x y] | s == "&&" = NaryOp And [x, y] + | s == "||" = NaryOp Or [x, y] | s == "<|>" = Msum [x, y] | otherwise = app2 s x y binary :: String -> Ex.Assoc -> Operator (AST' Text) binary s = Ex.Infix (binop (Text.pack s) <$ reservedOp s) -prefix :: String -> (a -> a) -> Operator a +prefix :: String -> (a -> a) -> Operator a prefix s f = Ex.Prefix (f <$ reservedOp s) postfix :: Parser (a -> a) -> Operator a @@ -199,7 +201,7 @@ , [ binary "+" Ex.AssocLeft , binary "-" Ex.AssocLeft , prefix "-" (app1 "negate")] - -- TODO: add "<=", ">=", "/=" + -- TODO: add "/=" -- TODO: do you *really* mean AssocLeft? Shouldn't they be non-assoc? , [ postfix ann_expr ] , [ binary "<|>" Ex.AssocRight] @@ -208,7 +210,8 @@ , binary "<=" Ex.AssocLeft , binary ">=" Ex.AssocLeft , binary "==" Ex.AssocLeft] - , [ binary "&&" Ex.AssocLeft]] + , [ binary "&&" Ex.AssocLeft] + , [ binary "||" Ex.AssocLeft]] unit_ :: Parser (AST' a) unit_ = Unit <$ symbol "()" @@ -314,10 +317,10 @@ reserved "integrate" *> (Integrate <$> identifier - <* symbol "from" + <* symbol "from" <*> expr <* symbol "to" - <*> expr + <*> expr <*> semiblockExpr ) @@ -326,10 +329,10 @@ reserved "summate" *> (Summate <$> identifier - <* symbol "from" + <* symbol "from" <*> expr <* symbol "to" - <*> expr + <*> expr <*> semiblockExpr ) @@ -338,10 +341,10 @@ reserved "product" *> (Product <$> identifier - <* symbol "from" + <* symbol "from" <*> expr <* symbol "to" - <*> expr + <*> expr <*> semiblockExpr ) @@ -378,15 +381,7 @@ array_literal :: Parser (AST' Text) array_literal = checkEmpty <$> brackets (commaSep expr) where checkEmpty [] = Empty - checkEmpty xs = Array "" (ULiteral . Nat . fromIntegral . length $ xs) - (go 0 xs) - - go _ [] = error "the impossible happened" - go _ [x] = x - go n (x:xs) = If (Var "equal" `App` (Var "") `App` (ULiteral $ Nat n)) - x - (go (n + 1) xs) - + checkEmpty xs = ArrayLiteral xs plate_expr :: Parser (AST' Text) plate_expr = @@ -474,7 +469,7 @@ <|> try lam_expr <|> try def_expr <|> try match_expr - -- <|> try data_expr + <|> try data_expr <|> try integrate_expr <|> try summate_expr <|> try product_expr @@ -506,14 +501,15 @@ mkIndentStream 0 infIndentation True Ge . mkCharIndentStream parseHakaru :: Text -> Either ParseError (AST' Text) -parseHakaru = - runParser (skipMany (comments <|> emptyLine) *> - expr <* eof) () "" . indentConfig +parseHakaru = parseAtTopLevel expr parseHakaruWithImports :: Text -> Either ParseError (ASTWithImport' Text) -parseHakaruWithImports = +parseHakaruWithImports = parseAtTopLevel exprWithImport + +parseAtTopLevel :: Parser a -> Text -> Either ParseError a +parseAtTopLevel p = runParser (skipMany (comments <|> emptyLine) *> - exprWithImport <* eof) () "" . indentConfig + p <* eof) () "" . indentConfig . Text.strip withPos :: Parser (AST' a) -> Parser (AST' a) withPos x = do @@ -522,14 +518,26 @@ e <- getPosition return $ WithMeta x' (SourceSpan s e) +{- +user-defined types: + +data either(a,b): + left(a) + right(a) + +data maybe(a): + nothing + just(a) +-} + data_expr :: Parser (AST' Text) -data_expr = +data_expr = do reserved "data" - *> (Data - <$> identifier - <* parens (commaSep identifier) -- TODO: why throw them away? - <*> blockOfMany (try type_app <|> type_var) - ) + ident <- identifier + typvars <- parens (commaSep identifier) + ts <- blockOfMany (try type_app <|> type_var) + e <- expr + return (Data ident typvars ts e) import_expr :: Parser (Import Text) import_expr = diff --git a/haskell/Language/Hakaru/Parser/SymbolResolve.hs b/haskell/Language/Hakaru/Parser/SymbolResolve.hs --- a/haskell/Language/Hakaru/Parser/SymbolResolve.hs +++ b/haskell/Language/Hakaru/Parser/SymbolResolve.hs @@ -67,7 +67,7 @@ U.PKonst a `U.PEt` U.PKonst b `U.PEt` U.PDone primTypes :: [(Text, Symbol' U.SSing)] -primTypes = +primTypes = [ ("nat", TNeu' $ U.SSing SNat) , ("int", TNeu' $ U.SSing SInt) , ("prob", TNeu' $ U.SSing SProb) @@ -115,7 +115,6 @@ ,("gamma", primMeasure2 (U.SomeOp T.Gamma)) ,("beta", primMeasure2 (U.SomeOp T.Beta)) ,("categorical", primMeasure1 (U.SomeOp T.Categorical)) - ,("bern", primBern) ,("factor", primFactor) ,("weight", primWeight) ,("dirac", TLam $ TNeu . syn . U.Dirac_) @@ -225,19 +224,10 @@ TNeu . syn . U.Datum_ . U.Datum "nothing" . U.Inl $ U.Done -primWeight, primFactor, primBern :: Symbol U.AST +primWeight, primFactor :: Symbol U.AST primWeight = t2 $ \w m -> syn $ U.Superpose_ (singleton (w, m)) primFactor = TLam $ \w -> TNeu . syn . U.Superpose_ $ singleton (w, syn $ U.Dirac_ unit_) -primBern = - TLam $ \p -> TNeu . syn . U.Superpose_ . L.fromList $ - [ (p, syn $ U.Dirac_ true_) - , (unsafeFrom_ . syn $ U.NaryOp_ U.Sum - [ syn $ U.Literal_ (Some1 $ T.LReal 1.0) - , syn $ U.PrimOp_ U.Negate [p] - ] - , syn $ U.Dirac_ false_) - ] two :: U.AST two = syn . U.Literal_ . U.val . U.Nat $ 2 @@ -270,9 +260,9 @@ name' <- gensym name f (mkSym name') <$> symbolResolution symbols e1 - <*> symbolResolution (insertSymbol name' symbols) e2 - + <*> symbolResolution (insertSymbol name' symbols) e2 + -- TODO: clean up by merging the @Reader (SymbolTable)@ and @State Int@ monads -- | Figure out symbols and types. symbolResolution @@ -305,27 +295,34 @@ U.Infinity' -> return $ U.Infinity' U.ULiteral v -> return $ U.ULiteral v - U.Integrate name e1 e2 e3 -> do + U.Integrate name e1 e2 e3 -> do name' <- gensym name U.Integrate (mkSym name') <$> symbolResolution symbols e1 <*> symbolResolution symbols e2 - <*> symbolResolution (insertSymbol name' symbols) e3 + <*> symbolResolution (insertSymbol name' symbols) e3 - U.Summate name e1 e2 e3 -> do + U.Summate name e1 e2 e3 -> do name' <- gensym name U.Summate (mkSym name') <$> symbolResolution symbols e1 <*> symbolResolution symbols e2 - <*> symbolResolution (insertSymbol name' symbols) e3 + <*> symbolResolution (insertSymbol name' symbols) e3 - U.Product name e1 e2 e3 -> do + U.Product name e1 e2 e3 -> do name' <- gensym name U.Product (mkSym name') <$> symbolResolution symbols e1 <*> symbolResolution symbols e2 - <*> symbolResolution (insertSymbol name' symbols) e3 + <*> symbolResolution (insertSymbol name' symbols) e3 + U.Bucket name e1 e2 e3 -> do + name' <- gensym name + U.Bucket (mkSym name') + <$> symbolResolution symbols e1 + <*> symbolResolution symbols e2 + <*> symbolResolutionReducer (insertSymbol name' symbols) e3 + U.NaryOp op es -> U.NaryOp op <$> mapM (symbolResolution symbols) es @@ -337,6 +334,8 @@ U.Array name e1 e2 -> resolveBinder symbols name e1 e2 U.Array + U.ArrayLiteral es -> U.ArrayLiteral <$> mapM (symbolResolution symbols) es + U.Index a i -> U.Index <$> symbolResolution symbols a <*> symbolResolution symbols i @@ -349,24 +348,49 @@ U.Bind name e1 e2 -> resolveBinder symbols name e1 e2 U.Bind U.Plate name e1 e2 -> resolveBinder symbols name e1 e2 U.Plate U.Expect name e1 e2 -> resolveBinder symbols name e1 e2 U.Expect - U.Chain name e1 e2 e3 -> do + U.Chain name e1 e2 e3 -> do name' <- gensym name U.Chain (mkSym name') <$> symbolResolution symbols e1 <*> symbolResolution symbols e2 - <*> symbolResolution (insertSymbol name' symbols) e3 + <*> symbolResolution (insertSymbol name' symbols) e3 U.Observe e1 e2 -> U.Observe <$> symbolResolution symbols e1 <*> symbolResolution symbols e2 U.Msum es -> U.Msum <$> mapM (symbolResolution symbols) es - U.Data _name _typ -> error "TODO: symbolResolution{U.Data}" + U.Data name tvars typ e -> error $ ("TODO: symbolResolution{U.Data} " ++ + show name ++ " with " ++ + show tvars ++ ":" ++ show typ) U.WithMeta a meta -> U.WithMeta <$> symbolResolution symbols a <*> return meta +symbolResolutionReducer + :: SymbolTable + -> U.Reducer' Text + -> State Int (U.Reducer' (Symbol U.AST)) +symbolResolutionReducer symbols ast = + case ast of + U.R_Fanout e1 e2 -> U.R_Fanout + <$> symbolResolutionReducer symbols e1 + <*> symbolResolutionReducer symbols e2 + U.R_Index name e1 e2 e3 -> do + name' <- gensym name + U.R_Index (mkSym name') + <$> symbolResolution symbols e1 + <*> symbolResolution symbols e2 + <*> symbolResolutionReducer (insertSymbol name' symbols) e3 + U.R_Split e1 e2 e3 -> U.R_Split + <$> symbolResolution symbols e1 + <*> symbolResolutionReducer symbols e2 + <*> symbolResolutionReducer symbols e3 + U.R_Nop -> return U.R_Nop + U.R_Add e1 -> U.R_Add <$> symbolResolution symbols e1 + + symbolResolveBranch :: SymbolTable -> U.Branch' Text @@ -420,13 +444,15 @@ U.Integrate name e1 e2 e3 -> U.Integrate name (normAST e1) (normAST e2) (normAST e3) U.Summate name e1 e2 e3 -> U.Summate name (normAST e1) (normAST e2) (normAST e3) U.Product name e1 e2 e3 -> U.Product name (normAST e1) (normAST e2) (normAST e3) + U.Bucket name e1 e2 e3 -> U.Bucket name (normAST e1) (normAST e2) (redNorm e3) U.ULiteral v -> U.ULiteral v U.NaryOp op es -> U.NaryOp op (map normAST es) U.Unit -> U.Unit U.Empty -> U.Empty U.Pair e1 e2 -> U.Pair (normAST e1) (normAST e2) U.Array name e1 e2 -> U.Array name (normAST e1) (normAST e2) - U.Index e1 e2 -> U.Index (normAST e1) (normAST e2) + U.ArrayLiteral es -> U.ArrayLiteral (map normAST es) + U.Index e1 e2 -> U.Index (normAST e1) (normAST e2) U.Case e1 e2 -> U.Case (normAST e1) (map branchNorm e2) U.Dirac e1 -> U.Dirac (normAST e1) U.Bind name e1 e2 -> U.Bind name (normAST e1) (normAST e2) @@ -435,13 +461,26 @@ U.Expect name e1 e2 -> U.Expect name (normAST e1) (normAST e2) U.Observe e1 e2 -> U.Observe (normAST e1) (normAST e2) U.Msum es -> U.Msum (map normAST es) - U.Data name typ -> U.Data name typ + U.Data name tvars typs e -> U.Data name tvars typs e + -- do we need to norm here? what if we try to define `true` which is already a constructor U.WithMeta a meta -> U.WithMeta (normAST a) meta branchNorm :: U.Branch' (Symbol U.AST) -> U.Branch' (Symbol U.AST) branchNorm (U.Branch' pat e2') = U.Branch' pat (normAST e2') branchNorm (U.Branch'' pat e2') = U.Branch'' pat (normAST e2') +redNorm :: U.Reducer' (Symbol U.AST) -> U.Reducer' (Symbol U.AST) +redNorm ast = + case ast of + U.R_Fanout e1 e2 -> + U.R_Fanout (redNorm e1) (redNorm e2) + U.R_Index name e1 e2 e3 -> + U.R_Index name (normAST e1) (normAST e2) (redNorm e3) + U.R_Split e1 e2 e3 -> + U.R_Split (normAST e1) (redNorm e2) (redNorm e3) + U.R_Nop -> U.R_Nop + U.R_Add e1 -> U.R_Add (normAST e1) + collapseSuperposes :: [U.AST] -> U.AST collapseSuperposes es = syn $ U.Superpose_ (fromList $ F.concatMap go es) where @@ -490,8 +529,30 @@ U.Branch_ (makePattern (U.PData' (U.DV "true" []))) (makeAST e) makeFalse e = U.Branch_ (makePattern (U.PData' (U.DV "false" []))) (makeAST e) - - + +makeReducerAST + :: Variable 'U.U + -> U.Reducer' (Symbol U.AST) + -> List1 Variable xs + -> U.Reducer xs U.U_ABT 'U.U +makeReducerAST i r1 bs = + case r1 of + U.R_Fanout r2 r3 -> U.R_Fanout_ + (makeReducerAST i r2 bs) + (makeReducerAST i r3 bs) + U.R_Index b e1 e2 r1 -> withName "U.R_Index" b $ \b' -> + U.R_Index_ + b' -- HACK: This shouldn't be needed here + (binds_ bs (makeAST e1)) + (bind i (binds_ bs (makeAST e2))) + (makeReducerAST i r1 (Cons1 b' bs)) + U.R_Split e1 r2 r3 -> U.R_Split_ + (bind i (binds_ bs (makeAST e1))) + (makeReducerAST i r2 bs) + (makeReducerAST i r3 bs) + U.R_Nop -> U.R_Nop_ + U.R_Add e1 -> U.R_Add_ (bind i (binds_ bs (makeAST e1))) + makeAST :: U.AST' (Symbol U.AST) -> U.AST makeAST ast = case ast of @@ -520,6 +581,7 @@ U.Array s e1 e2 -> withName "U.Array" s $ \name -> syn $ U.Array_ (makeAST e1) (bind name $ makeAST e2) + U.ArrayLiteral es -> syn $ U.ArrayLiteral_ (map makeAST es) U.Index e1 e2 -> syn $ U.ArrayOp_ U.Index_ [(makeAST e1), (makeAST e2)] U.Case e bs -> syn $ U.Case_ (makeAST e) (map makeBranch bs) U.Dirac e1 -> syn $ U.Dirac_ (makeAST e1) @@ -541,21 +603,23 @@ U.Product s e1 e2 e3 -> withName "U.Product" s $ \name -> syn $ U.Product_ (makeAST e1) (makeAST e2) (bind name $ makeAST e3) + U.Bucket s e1 e2 e3 -> + withName "U.Bucket" s $ \name -> + syn $ U.Bucket_ (makeAST e1) (makeAST e2) (makeReducerAST name e3 Nil1) U.Expect s e1 e2 -> withName "U.Expect" s $ \name -> syn $ U.Expect_ (makeAST e1) (bind name $ makeAST e2) U.Observe e1 e2 -> syn $ U.Observe_ (makeAST e1) (makeAST e2) U.Msum es -> collapseSuperposes (map makeAST es) - - U.Data _name _typ -> error "TODO: makeAST{U.Data}" + + U.Data name tvars typs e -> error "TODO: makeAST{U.Data}" U.WithMeta a meta -> withMetadata meta (makeAST a) - where - withName :: String -> Symbol U.AST -> (Variable 'U.U -> r) -> r - withName fun s k = - case s of - TNeu e -> caseVarSyn e k (error $ "makeAST: bad " ++ fun) - _ -> error $ "makeAST: bad " ++ fun +withName :: String -> Symbol U.AST -> (Variable 'U.U -> r) -> r +withName fun s k = + case s of + TNeu e -> caseVarSyn e k (error $ "makeAST: bad " ++ fun) + _ -> error $ "makeAST: bad " ++ fun resolveAST :: U.AST' Text -> U.AST resolveAST ast = diff --git a/haskell/Language/Hakaru/Pretty/Concrete.hs b/haskell/Language/Hakaru/Pretty/Concrete.hs --- a/haskell/Language/Hakaru/Pretty/Concrete.hs +++ b/haskell/Language/Hakaru/Pretty/Concrete.hs @@ -121,8 +121,8 @@ go xs (Syn t) = (reverse xs, prettyPrec_ 0 (LC_ (syn t))) -ppBinder2 :: (ABT Term abt) => abt xs a -> ([Doc], [Doc], Docs) -ppBinder2 e = unpackVarTypes $ go [] (viewABT e) +ppBinder2prec :: (ABT Term abt) => Int -> abt xs a -> ([Doc], [Doc], Docs) +ppBinder2prec p e = unpackVarTypes $ go [] (viewABT e) where unpackVarTypes (varTypes, body) = (map fst varTypes, map snd varTypes, body) @@ -130,8 +130,10 @@ go :: (ABT Term abt) => [(Doc, Doc)] -> View (Term abt) xs a -> ([(Doc, Doc)],Docs) go xs (Bind x v) = go (ppVariableWithType x : xs) v go xs (Var x) = (reverse xs, [ppVariable x]) - go xs (Syn t) = (reverse xs, prettyPrec_ 0 (LC_ (syn t))) + go xs (Syn t) = (reverse xs, prettyPrec_ p (LC_ (syn t))) +ppBinder2 :: (ABT Term abt) => abt xs a -> ([Doc], [Doc], Docs) +ppBinder2 = ppBinder2prec 0 -- TODO: since switching to ABT2, this instance requires -XFlexibleContexts; we should fix that if we can -- BUG: since switching to ABT2, this instance requires -XUndecidableInstances; must be fixed! @@ -181,8 +183,8 @@ identityElement Or = [PP.text "false"] identityElement Xor = [PP.text "false"] identityElement Iff = [PP.text "true"] - identityElement (Min _) = error "min can not be used with no arguements" - identityElement (Max _) = error "max can not be used with no arguements" + identityElement (Min _) = error "min cannot be used with no arguments" + identityElement (Max _) = error "max cannot be used with no arguments" identityElement (Sum _) = [PP.text "0"] identityElement (Prod _) = [PP.text "1"] @@ -197,6 +199,8 @@ <> PP.colon <> PP.space , PP.nest 1 (toDoc body)] + ArrayLiteral_ es -> ppList $ fmap (prettyPrec 0) es + Datum_ d -> prettyPrec_ p (fmap11 LC_ d) Case_ e1 bs -> parens True $ -- TODO: should we also add hints to the 'Case_' constructor to know whether it came from 'if_', 'unpair', etc? @@ -255,8 +259,9 @@ -- | Pretty-print @(:$)@ nodes in the AST. ppSCon :: (ABT Term abt) => Int -> SCon args a -> SArgs abt args -> Docs -ppSCon _ Lam_ = \(e1 :* End) -> - let (vars, types, body) = ppBinder2 e1 in +ppSCon p Lam_ = \(e1 :* End) -> + let (vars, types, body) = ppBinder2prec 11 e1 in + parens (p < 11) [ PP.text "fn" <+> toDoc vars <+> toDoc types <> PP.colon <> PP.space @@ -321,6 +326,7 @@ ppSCon p (Summate _ _) = \(e1 :* e2 :* e3 :* End) -> let (vars, types, body) = ppBinder2 e3 in + parens True $ [ PP.text "summate" <+> toDoc vars <+> PP.text "from" @@ -333,6 +339,7 @@ ppSCon p (Product _ _) = \(e1 :* e2 :* e3 :* End) -> let (vars, types, body) = ppBinder2 e3 in + parens True $ [ PP.text "product" <+> toDoc vars <+> PP.text "from" @@ -393,7 +400,8 @@ prettyType _ SReal = PP.text "real" prettyType p (SMeasure a) = PP.text "measure" <> PP.parens (prettyType p a) prettyType p (SArray a) = PP.text "array" <> PP.parens (prettyType p a) -prettyType p (SFun a b) = prettyType p a <+> PP.text "->" <+> prettyType p b +-- HACK: precedence of function types +prettyType p (SFun a b) = PP.parens (prettyType p a <+> PP.text "->" <+> prettyType p b) prettyType p typ = case typ of SData (STyCon sym `STyApp` a `STyApp` b) _ diff --git a/haskell/Language/Hakaru/Pretty/Haskell.hs b/haskell/Language/Hakaru/Pretty/Haskell.hs --- a/haskell/Language/Hakaru/Pretty/Haskell.hs +++ b/haskell/Language/Hakaru/Pretty/Haskell.hs @@ -53,6 +53,7 @@ import Language.Hakaru.Types.HClasses import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.ABT ---------------------------------------------------------------- @@ -111,10 +112,22 @@ -- | Pretty-print Hakaru binders as a Haskell lambda, as per our HOAS API. ppBinder :: (ABT Term abt) => abt xs a -> Docs ppBinder e = - case go [] (viewABT e) of + case ppViewABT e of ([], body) -> body (vars,body) -> PP.char '\\' <+> PP.sep vars <+> PP.text "->" : body + +ppUncurryBinder :: (ABT Term abt) => abt xs a -> Docs +ppUncurryBinder e = + case ppViewABT e of + (vars,body) -> PP.char '\\' <+> unc vars <+> PP.text "->" : body where + unc :: [Doc] -> Doc + unc [] = PP.text "()" + unc (x:xs) = PP.parens (x <> PP.comma <> unc xs) + +ppViewABT :: (ABT Term abt) => abt xs a -> ([Doc], Docs) +ppViewABT e = go [] (viewABT e) + where go :: (ABT Term abt) => [Doc] -> View (Term abt) xs a -> ([Doc],Docs) go xs (Syn t) = (reverse xs, prettyPrec_ 0 (LC_ (syn t))) go xs (Var x) = (reverse xs, [ppVariable x]) @@ -171,6 +184,7 @@ [ ppArg e1 <+> PP.char '$' , toDoc $ ppBinder e2 ] + ArrayLiteral_ es -> ppFun 11 "arrayLit" (ppList $ map (prettyPrec 0) es) Datum_ d -> prettyPrec_ p (fmap11 LC_ d) Case_ e1 bs -> -- TODO: should we also add hints to the 'Case_' constructor to know whether it came from 'if_', 'unpair', etc? @@ -178,6 +192,13 @@ [ ppArg e1 , toDoc $ ppList (map (toDoc . prettyPrec_ 0) bs) ] + Bucket b e r -> + ppFun p "bucket" + [ ppArg b + , ppArg e + , toDoc $ parens True (prettyPrec_ p r) + ] + Superpose_ pes -> case pes of (e1,e2) L.:| [] -> @@ -426,6 +447,30 @@ -- have them decide if they need to or not. ] +instance (ABT Term abt) => Pretty (Reducer abt xs) where + prettyPrec_ p (Red_Fanout r1 r2) = + ppFun p "r_fanout" + [ toDoc $ prettyPrec_ 11 r1 + , toDoc $ prettyPrec_ 11 r2 + ] + prettyPrec_ p (Red_Index n o e) = + ppFun p "r_index" + [ toDoc $ parens True $ ppUncurryBinder n + , toDoc $ parens True $ ppUncurryBinder o + , toDoc $ prettyPrec_ 11 e + ] + prettyPrec_ p (Red_Split b r1 r2) = + ppFun p "r_split" + [ toDoc $ parens True (ppUncurryBinder b) + , toDoc $ prettyPrec_ 11 r1 + , toDoc $ prettyPrec_ 11 r2 + ] + prettyPrec_ p Red_Nop = + [ PP.text "r_nop" ] + prettyPrec_ p (Red_Add _ e) = + ppFun p "r_add" + [ toDoc $ parens True (ppUncurryBinder e)] + ---------------------------------------------------------------- -- | For the \"@lam $ \x ->\n@\" style layout. adjustHead :: (Doc -> Doc) -> Docs -> Docs diff --git a/haskell/Language/Hakaru/Pretty/Maple.hs b/haskell/Language/Hakaru/Pretty/Maple.hs --- a/haskell/Language/Hakaru/Pretty/Maple.hs +++ b/haskell/Language/Hakaru/Pretty/Maple.hs @@ -67,7 +67,6 @@ app3 fn x y z = op3 fn (arg x) (arg y) (arg z) {-# INLINE app3 #-} --- HACK: why doesn't Foldable imply Functor anymore? appN :: (ABT Term abt, Functor f, F.Foldable f) => String -> f (abt '[] a) -> ShowS appN fn xs = opN fn (arg <$> xs) @@ -105,6 +104,10 @@ parens a = showChar '(' . a . showChar ')' {-# INLINE parens #-} +brackets :: ShowS -> ShowS +brackets a = showChar '[' . a . showChar ']' +{-# INLINE brackets #-} + intercalate :: F.Foldable f => ShowS -> f ShowS -> ShowS intercalate sep = F.foldr1 (\a b -> a . sep . b) {-# INLINE intercalate #-} @@ -117,21 +120,22 @@ mapleAST (LC_ e) = caseVarSyn e var1 $ \t -> case t of - o :$ es -> mapleSCon o es - NaryOp_ op es -> mapleNary op es - Literal_ v -> mapleLiteral v - Empty_ _ -> error "TODO: mapleAST{Empty}" - Array_ e1 e2 -> + o :$ es -> mapleSCon o es + NaryOp_ op es -> mapleNary op es + Literal_ v -> mapleLiteral v + Empty_ _ -> error "TODO: mapleAST{Empty}" + Array_ e1 e2 -> caseBind e2 $ \x e2' -> app3 "ary" e1 (var x) e2' + ArrayLiteral_ es -> brackets (commaSep (fmap arg es)) Datum_ (Datum "true" _typ (Inl Done) ) -> showString "true" Datum_ (Datum "false" _typ (Inr (Inl Done))) -> showString "false" - Datum_ d -> mapleDatum d - Case_ e' bs -> + Datum_ d -> mapleDatum d + Case_ e' bs -> op2 "case" (arg e') (opN "Branches" (mapleBranch <$> bs)) - Superpose_ pms -> + Superpose_ pms -> opN "Msum" (uncurry (app2 "Weight") <$> L.toList pms) - Reject_ _ -> showString "Msum()" + Reject_ _ -> showString "Msum()" mapleLiteral :: Literal a -> ShowS @@ -181,7 +185,7 @@ error "TODO: mapleSCon{Chain}" mapleSCon Integrate = \(e1 :* e2 :* e3 :* End) -> caseBind e3 $ \x e3' -> - showString "int(" + showString "Int(" . arg e3' . showString ", [" . var1 x @@ -192,7 +196,7 @@ . showString "])" mapleSCon (Summate _ _) = \(e1 :* e2 :* e3 :* End) -> caseBind e3 $ \x e3' -> - showString "sum(" + showString "Sum(" . arg e3' . showString ", " . var1 x @@ -203,7 +207,7 @@ . showString ")-1)" mapleSCon (Product _ _) = \(e1 :* e2 :* e3 :* End) -> caseBind e3 $ \x e3' -> - showString "product(" + showString "Product(" . arg e3' . showString ", " . var1 x @@ -225,6 +229,7 @@ mapleNary :: (ABT Term abt) => NaryOp a -> Seq (abt '[] a) -> ShowS mapleNary And = appN "And" +mapleNary Or = appN "Or" mapleNary (Sum _) = parens . intercalate (showString " + ") . fmap arg mapleNary (Prod _) = parens . intercalate (showString " * ") . fmap arg mapleNary (Min _) = appN "min" @@ -302,6 +307,7 @@ maplePrimOp Not (e1 :* End) = app1 "Not" e1 maplePrimOp Pi End = showString "Pi" maplePrimOp Cos (e1 :* End) = app1 "cos" e1 +maplePrimOp Sin (e1 :* End) = app1 "sin" e1 maplePrimOp RealPow (e1 :* e2 :* End) = parens (arg e1 . showString " ^ " . arg e2) maplePrimOp Exp (e1 :* End) = app1 "exp" e1 @@ -332,8 +338,8 @@ mapleMeasureOp :: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) => MeasureOp typs a -> SArgs abt args -> ShowS -mapleMeasureOp Lebesgue = \End -> showString "Lebesgue()" -mapleMeasureOp Counting = \End -> showString "Counting()" +mapleMeasureOp Lebesgue = \End -> showString "Lebesgue(-infinity,infinity)" +mapleMeasureOp Counting = \End -> showString "Counting(-infinity,infinity)" mapleMeasureOp Categorical = \(e1 :* End) -> app1 "Categorical" e1 mapleMeasureOp Uniform = \(e1 :* e2 :* End) -> app2 "Uniform" e1 e2 mapleMeasureOp Normal = \(e1 :* e2 :* End) -> app2 "Gaussian" e1 e2 diff --git a/haskell/Language/Hakaru/Runtime/CmdLine.hs b/haskell/Language/Hakaru/Runtime/CmdLine.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Runtime/CmdLine.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE CPP, + FlexibleContexts, + FlexibleInstances, + UndecidableInstances, + TypeFamilies #-} +module Language.Hakaru.Runtime.CmdLine where + +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Generic as G +import qualified System.Random.MWC as MWC +import Language.Hakaru.Runtime.Prelude +import Data.Number.LogFloat +import Control.Monad (forever) + +#if __GLASGOW_HASKELL__ < 710 +import Data.Functor +#endif + +-- A class of types that can be parsed from command line arguments +class Parseable a where + parse :: String -> IO a + +instance Parseable Int where + parse = return . read + +instance Parseable Double where + parse = return . read + +instance (U.Unbox a, Parseable a) => Parseable (U.Vector a) where + parse s = U.fromList <$> ((mapM parse) =<< (lines <$> readFile s)) + +{- Make main needs to recur down the function type while at the term level build +-- up a continuation of parses and partial application of the function +-} +class MakeMain p where + makeMain :: p -> [String] -> IO () + +instance MakeMain Int where + makeMain p _ = print p + +instance MakeMain Double where + makeMain p _ = print p + +instance Show a => MakeMain (Measure a) where + makeMain p _ = MWC.createSystemRandom >>= \gen -> + forever $ do + ms <- unMeasure p gen + case ms of + Nothing -> return () + Just s -> print s + +instance {-# OVERLAPPABLE #-} + ( Show (v a), (G.Vector (MayBoxVec a) a), v ~ MayBoxVec a) + => MakeMain (v a) where + makeMain p _ = print p + +instance {-# OVERLAPPING #-}(Parseable a, MakeMain b) + => MakeMain (a -> b) where + makeMain p (a:as) = do a' <- parse a + makeMain (p a') as + makeMain _ [] = error "not enough arguments" diff --git a/haskell/Language/Hakaru/Runtime/LogFloatCmdLine.hs b/haskell/Language/Hakaru/Runtime/LogFloatCmdLine.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Runtime/LogFloatCmdLine.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE CPP, + FlexibleContexts, + FlexibleInstances, + UndecidableInstances, + TypeFamilies #-} +module Language.Hakaru.Runtime.LogFloatCmdLine where + +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Generic as G +import qualified System.Random.MWC as MWC +import Language.Hakaru.Runtime.LogFloatPrelude +import Data.Number.LogFloat +import Control.Monad (forever) + +#if __GLASGOW_HASKELL__ < 710 +import Data.Functor +#endif + +-- This Read instance really should be the logfloat package +instance Read LogFloat where + readsPrec p s = [(logFloat x, r) | (x, r) <- readsPrec p s] + +-- A class of types that can be parsed from command line arguments +class Parseable a where + parse :: String -> IO a + +instance Parseable Int where + parse = return . read + +instance Parseable Double where + parse = return . read + +instance Parseable LogFloat where + parse = return . read + +instance (U.Unbox a, Parseable a) => Parseable (U.Vector a) where + parse s = U.fromList <$> ((mapM parse) =<< (lines <$> readFile s)) + +instance (Read a, Read b, Parseable a, Parseable b) => Parseable (a, b) where + parse = return . read + +{- Make main needs to recur down the function type while at the term level build +-- up a continuation of parses and partial application of the function +-} +class MakeMain p where + makeMain :: p -> [String] -> IO () + +instance MakeMain Int where + makeMain p _ = print p + +instance MakeMain Double where + makeMain p _ = print p + +instance MakeMain LogFloat where + makeMain p _ = print p + +instance Show a => MakeMain (Measure a) where + makeMain p _ = MWC.createSystemRandom >>= \gen -> + forever $ do + ms <- unMeasure p gen + case ms of + Nothing -> return () + Just s -> print s + +instance {-# OVERLAPPABLE #-} + ( Show (v a), (G.Vector (MayBoxVec a) a), v ~ MayBoxVec a) + => MakeMain (v a) where + makeMain p _ = print p + +instance {-# OVERLAPPING #-}(Parseable a, MakeMain b) + => MakeMain (a -> b) where + makeMain p (a:as) = do a' <- parse a + makeMain (p a') as + makeMain _ [] = error "not enough arguments" diff --git a/haskell/Language/Hakaru/Runtime/LogFloatPrelude.hs b/haskell/Language/Hakaru/Runtime/LogFloatPrelude.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Runtime/LogFloatPrelude.hs @@ -0,0 +1,444 @@ +{-# LANGUAGE CPP + , GADTs + , Rank2Types + , DataKinds + , TypeFamilies + , FlexibleContexts + , UndecidableInstances + , LambdaCase + , MultiParamTypeClasses + , OverloadedStrings + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs -fsimpl-tick-factor=1000 #-} +module Language.Hakaru.Runtime.LogFloatPrelude where + +#if __GLASGOW_HASKELL__ < 710 +import Data.Functor ((<$>)) +import Control.Applicative (Applicative(..)) +#endif +import Data.Foldable as F +import qualified System.Random.MWC as MWC +import qualified System.Random.MWC.Distributions as MWCD +import Data.Number.Natural +import Data.Number.LogFloat hiding (sum, product) +import qualified Data.Number.LogFloat as LF +import Data.STRef +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Generic as G +import qualified Data.Vector.Generic.Mutable as M +import Control.Monad +import Control.Monad.ST +import Prelude hiding (init, sum, product, exp, log, (**), pi) +import qualified Prelude as P + +type family MinBoxVec (v1 :: * -> *) (v2 :: * -> *) :: * -> * +type instance MinBoxVec V.Vector v = V.Vector +type instance MinBoxVec v V.Vector = V.Vector +type instance MinBoxVec U.Vector U.Vector = U.Vector + +type family MayBoxVec a :: * -> * +type instance MayBoxVec () = U.Vector +type instance MayBoxVec Int = U.Vector +type instance MayBoxVec Double = U.Vector +type instance MayBoxVec LogFloat = U.Vector +type instance MayBoxVec Bool = U.Vector +type instance MayBoxVec (U.Vector a) = V.Vector +type instance MayBoxVec (V.Vector a) = V.Vector +type instance MayBoxVec (a,b) = MinBoxVec (MayBoxVec a) (MayBoxVec b) + +newtype instance U.MVector s LogFloat = MV_LogFloat (U.MVector s Double) +newtype instance U.Vector LogFloat = V_LogFloat (U.Vector Double) + +instance U.Unbox LogFloat + +instance M.MVector U.MVector LogFloat where + {-# INLINE basicLength #-} + {-# INLINE basicUnsafeSlice #-} + {-# INLINE basicOverlaps #-} + {-# INLINE basicUnsafeNew #-} +#if __GLASGOW_HASKELL__ > 710 + {-# INLINE basicInitialize #-} +#endif + {-# INLINE basicUnsafeReplicate #-} + {-# INLINE basicUnsafeRead #-} + {-# INLINE basicUnsafeWrite #-} + {-# INLINE basicClear #-} + {-# INLINE basicSet #-} + {-# INLINE basicUnsafeCopy #-} + {-# INLINE basicUnsafeGrow #-} + basicLength (MV_LogFloat v) = M.basicLength v + basicUnsafeSlice i n (MV_LogFloat v) = MV_LogFloat $ M.basicUnsafeSlice i n v + basicOverlaps (MV_LogFloat v1) (MV_LogFloat v2) = M.basicOverlaps v1 v2 + basicUnsafeNew n = MV_LogFloat `liftM` M.basicUnsafeNew n +#if __GLASGOW_HASKELL__ > 710 + basicInitialize (MV_LogFloat v) = M.basicInitialize v +#endif + basicUnsafeReplicate n x = MV_LogFloat `liftM` M.basicUnsafeReplicate n (logFromLogFloat x) + basicUnsafeRead (MV_LogFloat v) i = logToLogFloat `liftM` M.basicUnsafeRead v i + basicUnsafeWrite (MV_LogFloat v) i x = M.basicUnsafeWrite v i (logFromLogFloat x) + basicClear (MV_LogFloat v) = M.basicClear v + basicSet (MV_LogFloat v) x = M.basicSet v (logFromLogFloat x) + basicUnsafeCopy (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeCopy v1 v2 + basicUnsafeMove (MV_LogFloat v1) (MV_LogFloat v2) = M.basicUnsafeMove v1 v2 + basicUnsafeGrow (MV_LogFloat v) n = MV_LogFloat `liftM` M.basicUnsafeGrow v n + +instance G.Vector U.Vector LogFloat where + {-# INLINE basicUnsafeFreeze #-} + {-# INLINE basicUnsafeThaw #-} + {-# INLINE basicLength #-} + {-# INLINE basicUnsafeSlice #-} + {-# INLINE basicUnsafeIndexM #-} + {-# INLINE elemseq #-} + basicUnsafeFreeze (MV_LogFloat v) = V_LogFloat `liftM` G.basicUnsafeFreeze v + basicUnsafeThaw (V_LogFloat v) = MV_LogFloat `liftM` G.basicUnsafeThaw v + basicLength (V_LogFloat v) = G.basicLength v + basicUnsafeSlice i n (V_LogFloat v) = V_LogFloat $ G.basicUnsafeSlice i n v + basicUnsafeIndexM (V_LogFloat v) i + = logToLogFloat `liftM` G.basicUnsafeIndexM v i + basicUnsafeCopy (MV_LogFloat mv) (V_LogFloat v) + = G.basicUnsafeCopy mv v + elemseq _ x z = G.elemseq (undefined :: U.Vector a) (logFromLogFloat x) z + + +lam :: (a -> b) -> a -> b +lam = id +{-# INLINE lam #-} + +app :: (a -> b) -> a -> b +app f x = f x +{-# INLINE app #-} + +let_ :: a -> (a -> b) -> b +let_ x f = let x1 = x in f x1 +{-# INLINE let_ #-} + +ann_ :: a -> b -> b +ann_ _ a = a +{-# INLINE ann_ #-} + +exp :: Double -> LogFloat +exp = logToLogFloat +{-# INLINE exp #-} + +log :: LogFloat -> Double +log = logFromLogFloat +{-# INLINE log #-} + +newtype Measure a = Measure { unMeasure :: MWC.GenIO -> IO (Maybe a) } + +instance Functor Measure where + fmap = liftM + {-# INLINE fmap #-} + +instance Applicative Measure where + pure x = Measure $ \_ -> return (Just x) + {-# INLINE pure #-} + (<*>) = ap + {-# INLINE (<*>) #-} + +instance Monad Measure where + return = pure + {-# INLINE return #-} + m >>= f = Measure $ \g -> do + Just x <- unMeasure m g + unMeasure (f x) g + {-# INLINE (>>=) #-} + +makeMeasure :: (MWC.GenIO -> IO a) -> Measure a +makeMeasure f = Measure $ \g -> Just <$> f g +{-# INLINE makeMeasure #-} + +uniform :: Double -> Double -> Measure Double +uniform lo hi = makeMeasure $ MWC.uniformR (lo, hi) +{-# INLINE uniform #-} + +normal :: Double -> LogFloat -> Measure Double +normal mu sd = makeMeasure $ MWCD.normal mu (fromLogFloat sd) +{-# INLINE normal #-} + +beta :: LogFloat -> LogFloat -> Measure LogFloat +beta a b = makeMeasure $ \g -> + logFloat <$> MWCD.beta (fromLogFloat a) (fromLogFloat b) g +{-# INLINE beta #-} + +gamma :: LogFloat -> LogFloat -> Measure LogFloat +gamma a b = makeMeasure $ \g -> + logFloat <$> MWCD.gamma (fromLogFloat a) (fromLogFloat b) g +{-# INLINE gamma #-} + +categorical :: MayBoxVec LogFloat LogFloat -> Measure Int +categorical a = makeMeasure $ \g -> + fromIntegral <$> MWCD.categorical (U.map prep a) g + where prep p = fromLogFloat (p / m) + m = G.maximum a +{-# INLINE categorical #-} + +plate :: (G.Vector (MayBoxVec a) a) => + Int -> (Int -> Measure a) -> Measure (MayBoxVec a a) +plate n f = G.generateM (fromIntegral n) $ \x -> + f (fromIntegral x) +{-# INLINE plate #-} + +bucket :: Int -> Int -> (forall s. Reducer () s a) -> a +bucket b e r = runST + $ case r of Reducer{init=initR,accum=accumR,done=doneR} -> do + s' <- initR () + F.mapM_ (\i -> accumR () i s') [b .. e - 1] + doneR s' +{-# INLINE bucket #-} + +data Reducer xs s a = + forall cell. + Reducer { init :: xs -> ST s cell + , accum :: xs -> Int -> cell -> ST s () + , done :: cell -> ST s a + } + +r_fanout :: Reducer xs s a + -> Reducer xs s b + -> Reducer xs s (a,b) +r_fanout Reducer{init=initA,accum=accumA,done=doneA} + Reducer{init=initB,accum=accumB,done=doneB} = Reducer + { init = \xs -> liftM2 (,) (initA xs) (initB xs) + , accum = \bs i (s1, s2) -> + accumA bs i s1 >> accumB bs i s2 + , done = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2) + } +{-# INLINE r_fanout #-} + +r_index :: (G.Vector (MayBoxVec a) a) + => (xs -> Int) + -> ((Int, xs) -> Int) + -> Reducer (Int, xs) s a + -> Reducer xs s (MayBoxVec a a) +r_index n f Reducer{init=initR,accum=accumR,done=doneR} = Reducer + { init = \xs -> V.generateM (n xs) (\b -> initR (b, xs)) + , accum = \bs i v -> + let ov = f (i, bs) in + accumR (ov,bs) i (v V.! ov) + , done = \v -> fmap G.convert (V.mapM doneR v) + } +{-# INLINE r_index #-} + +r_split :: ((Int, xs) -> Bool) + -> Reducer xs s a + -> Reducer xs s b + -> Reducer xs s (a,b) +r_split b Reducer{init=initA,accum=accumA,done=doneA} + Reducer{init=initB,accum=accumB,done=doneB} = Reducer + { init = \xs -> liftM2 (,) (initA xs) (initB xs) + , accum = \bs i (s1, s2) -> + if (b (i,bs)) then accumA bs i s1 else accumB bs i s2 + , done = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2) + } +{-# INLINE r_split #-} + +r_add :: Num a => ((Int, xs) -> a) -> Reducer xs s a +r_add e = Reducer + { init = \_ -> newSTRef 0 + , accum = \bs i s -> + modifySTRef' s (+ (e (i,bs))) + , done = readSTRef + } +{-# INLINE r_add #-} + +r_nop :: Reducer xs s () +r_nop = Reducer + { init = \_ -> return () + , accum = \_ _ _ -> return () + , done = \_ -> return () + } +{-# INLINE r_nop #-} + +pair :: a -> b -> (a, b) +pair = (,) +{-# INLINE pair #-} + +true, false :: Bool +true = True +false = False + +nothing :: Maybe a +nothing = Nothing + +just :: a -> Maybe a +just = Just + +unit :: () +unit = () + +data Pattern = PVar | PWild +newtype Branch a b = + Branch { extract :: a -> Maybe b } + +ptrue, pfalse :: a -> Branch Bool a +ptrue b = Branch { extract = extractBool True b } +pfalse b = Branch { extract = extractBool False b } +{-# INLINE ptrue #-} +{-# INLINE pfalse #-} + +extractBool :: Bool -> a -> Bool -> Maybe a +extractBool b a p | p == b = Just a + | otherwise = Nothing +{-# INLINE extractBool #-} + +pnothing :: b -> Branch (Maybe a) b +pnothing b = Branch { extract = \ma -> case ma of + Nothing -> Just b + Just _ -> Nothing } + +pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b +pjust PVar c = Branch { extract = \ma -> case ma of + Nothing -> Nothing + Just x -> Just (c x) } +pjust _ _ = error "Runtime.Prelude pjust" + + +ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x,y) b +ppair PVar PVar c = Branch { extract = (\(x,y) -> Just (c x y)) } +ppair _ _ _ = error "ppair: TODO" + +uncase_ :: Maybe a -> a +uncase_ (Just a) = a +uncase_ Nothing = error "case_: unable to match any branches" +{-# INLINE uncase_ #-} + +case_ :: a -> [Branch a b] -> b +case_ e [c1] = uncase_ (extract c1 e) +case_ e [c1, c2] = uncase_ (extract c1 e `mplus` extract c2 e) +case_ e bs_ = go bs_ + where go [] = error "case_: unable to match any branches" + go (b:bs) = case extract b e of + Just b' -> b' + Nothing -> go bs +{-# INLINE case_ #-} + +branch :: (c -> Branch a b) -> c -> Branch a b +branch pat body = pat body +{-# INLINE branch #-} + +dirac :: a -> Measure a +dirac = return +{-# INLINE dirac #-} + +pose :: LogFloat -> Measure a -> Measure a +pose _ a = a +{-# INLINE pose #-} + +superpose :: [(LogFloat, Measure a)] + -> Measure a +superpose pms = do + i <- categorical (G.fromList $ map fst pms) + snd (pms !! i) +{-# INLINE superpose #-} + +reject :: Measure a +reject = Measure $ \_ -> return Nothing + +nat_ :: Int -> Int +nat_ = id + +int_ :: Int -> Int +int_ = id + +unsafeNat :: Int -> Int +unsafeNat = id + +nat2prob :: Int -> LogFloat +nat2prob = fromIntegral + +fromInt :: Int -> Double +fromInt = fromIntegral + +nat2int :: Int -> Int +nat2int = id + +nat2real :: Int -> Double +nat2real = fromIntegral + +fromProb :: LogFloat -> Double +fromProb = fromLogFloat + +unsafeProb :: Double -> LogFloat +unsafeProb = logFloat + +real_ :: Rational -> Double +real_ = fromRational + +prob_ :: NonNegativeRational -> LogFloat +prob_ = fromRational . fromNonNegativeRational + +infinity :: Double +infinity = 1/0 + +abs_ :: Num a => a -> a +abs_ = abs + +(**) :: LogFloat -> Double -> LogFloat +(**) = pow +{-# INLINE (**) #-} + +pi :: LogFloat +pi = logFloat P.pi +{-# INLINE pi #-} + +thRootOf :: Int -> LogFloat -> LogFloat +thRootOf a b = b `pow` (recip $ fromIntegral a) +{-# INLINE thRootOf #-} + +array + :: (G.Vector (MayBoxVec a) a) + => Int + -> (Int -> a) + -> MayBoxVec a a +array n f = G.generate (fromIntegral n) (f . fromIntegral) +{-# INLINE array #-} + +arrayLit :: (G.Vector (MayBoxVec a) a) => [a] -> MayBoxVec a a +arrayLit = G.fromList +{-# INLINE arrayLit #-} + +(!) :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int -> a +a ! b = a G.! (fromIntegral b) +{-# INLINE (!) #-} + +size :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int +size v = fromIntegral (G.length v) +{-# INLINE size #-} + +class Num a => Num' a where + product :: Int -> Int -> (Int -> a) -> a + product a b f = F.foldl' (\x y -> x * f y) 1 [a .. b-1] + {-# INLINE product #-} + summate :: Int -> Int -> (Int -> a) -> a + summate a b f = F.foldl' (\x y -> x + f y) 0 [a .. b-1] + {-# INLINE summate #-} + +instance Num' Int +instance Num' Double +instance Num' LogFloat where + product a b f = LF.product (map f [a .. b-1]) + {-# INLINE product #-} + summate a b f = LF.sum (map f [a .. b-1]) + {-# INLINE summate #-} + +run :: Show a + => MWC.GenIO + -> Measure a + -> IO () +run g k = unMeasure k g >>= \case + Just a -> print a + Nothing -> return () + +iterateM_ + :: Monad m + => (a -> m a) + -> a + -> m b +iterateM_ f = g + where g x = f x >>= g + +withPrint :: Show a => (a -> IO b) -> a -> IO b +withPrint f x = print x >> f x diff --git a/haskell/Language/Hakaru/Runtime/Prelude.hs b/haskell/Language/Hakaru/Runtime/Prelude.hs --- a/haskell/Language/Hakaru/Runtime/Prelude.hs +++ b/haskell/Language/Hakaru/Runtime/Prelude.hs @@ -6,6 +6,7 @@ , UndecidableInstances , LambdaCase , OverloadedStrings + , Rank2Types #-} {-# OPTIONS_GHC -Wall -fwarn-tabs -fsimpl-tick-factor=1000 #-} @@ -19,11 +20,13 @@ import qualified System.Random.MWC as MWC import qualified System.Random.MWC.Distributions as MWCD import Data.Number.Natural +import Data.STRef import qualified Data.Vector as V import qualified Data.Vector.Unboxed as U import qualified Data.Vector.Generic as G import Control.Monad -import Prelude hiding (product) +import Control.Monad.ST +import Prelude hiding (product, init) type family MinBoxVec (v1 :: * -> *) (v2 :: * -> *) :: * -> * type instance MinBoxVec V.Vector v = V.Vector @@ -31,69 +34,171 @@ type instance MinBoxVec U.Vector U.Vector = U.Vector type family MayBoxVec a :: * -> * +type instance MayBoxVec () = U.Vector type instance MayBoxVec Int = U.Vector type instance MayBoxVec Double = U.Vector +type instance MayBoxVec Bool = U.Vector type instance MayBoxVec (U.Vector a) = V.Vector type instance MayBoxVec (V.Vector a) = V.Vector type instance MayBoxVec (a,b) = MinBoxVec (MayBoxVec a) (MayBoxVec b) lam :: (a -> b) -> a -> b lam = id +{-# INLINE lam #-} app :: (a -> b) -> a -> b app f x = f x +{-# INLINE app #-} let_ :: a -> (a -> b) -> b let_ x f = let x1 = x in f x1 +{-# INLINE let_ #-} ann_ :: a -> b -> b ann_ _ a = a +{-# INLINE ann_ #-} newtype Measure a = Measure { unMeasure :: MWC.GenIO -> IO (Maybe a) } instance Functor Measure where fmap = liftM + {-# INLINE fmap #-} instance Applicative Measure where pure x = Measure $ \_ -> return (Just x) + {-# INLINE pure #-} (<*>) = ap + {-# INLINE (<*>) #-} instance Monad Measure where return = pure + {-# INLINE return #-} m >>= f = Measure $ \g -> do Just x <- unMeasure m g unMeasure (f x) g + {-# INLINE (>>=) #-} makeMeasure :: (MWC.GenIO -> IO a) -> Measure a makeMeasure f = Measure $ \g -> Just <$> f g +{-# INLINE makeMeasure #-} uniform :: Double -> Double -> Measure Double uniform lo hi = makeMeasure $ MWC.uniformR (lo, hi) +{-# INLINE uniform #-} normal :: Double -> Double -> Measure Double normal mu sd = makeMeasure $ MWCD.normal mu sd +{-# INLINE normal #-} beta :: Double -> Double -> Measure Double beta a b = makeMeasure $ MWCD.beta a b +{-# INLINE beta #-} gamma :: Double -> Double -> Measure Double gamma a b = makeMeasure $ MWCD.gamma a b +{-# INLINE gamma #-} categorical :: MayBoxVec Double Double -> Measure Int categorical a = makeMeasure (\g -> fromIntegral <$> MWCD.categorical a g) +{-# INLINE categorical #-} plate :: (G.Vector (MayBoxVec a) a) => Int -> (Int -> Measure a) -> Measure (MayBoxVec a a) plate n f = G.generateM (fromIntegral n) $ \x -> f (fromIntegral x) +{-# INLINE plate #-} +bucket :: Int -> Int -> (forall s. Reducer () s a) -> a +bucket b e r = runST + $ case r of Reducer{init=initR,accum=accumR,done=doneR} -> do + s' <- initR () + F.mapM_ (\i -> accumR () i s') [b .. e - 1] + doneR s' +{-# INLINE bucket #-} + +data Reducer xs s a = + forall cell. + Reducer { init :: xs -> ST s cell + , accum :: xs -> Int -> cell -> ST s () + , done :: cell -> ST s a + } + +r_fanout :: Reducer xs s a + -> Reducer xs s b + -> Reducer xs s (a,b) +r_fanout Reducer{init=initA,accum=accumA,done=doneA} + Reducer{init=initB,accum=accumB,done=doneB} = Reducer + { init = \xs -> liftM2 (,) (initA xs) (initB xs) + , accum = \bs i (s1, s2) -> + accumA bs i s1 >> accumB bs i s2 + , done = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2) + } +{-# INLINE r_fanout #-} + +r_index :: (G.Vector (MayBoxVec a) a) + => (xs -> Int) + -> ((Int, xs) -> Int) + -> Reducer (Int, xs) s a + -> Reducer xs s (MayBoxVec a a) +r_index n f Reducer{init=initR,accum=accumR,done=doneR} = Reducer + { init = \xs -> V.generateM (n xs) (\b -> initR (b, xs)) + , accum = \bs i v -> + let ov = f (i, bs) in + accumR (ov,bs) i (v V.! ov) + , done = \v -> fmap G.convert (V.mapM doneR v) + } +{-# INLINE r_index #-} + +r_split :: ((Int, xs) -> Bool) + -> Reducer xs s a + -> Reducer xs s b + -> Reducer xs s (a,b) +r_split b Reducer{init=initA,accum=accumA,done=doneA} + Reducer{init=initB,accum=accumB,done=doneB} = Reducer + { init = \xs -> liftM2 (,) (initA xs) (initB xs) + , accum = \bs i (s1, s2) -> + if (b (i,bs)) then accumA bs i s1 else accumB bs i s2 + , done = \(s1, s2) -> liftM2 (,) (doneA s1) (doneB s2) + } +{-# INLINE r_split #-} + +r_add :: Num a => ((Int, xs) -> a) -> Reducer xs s a +r_add e = Reducer + { init = \_ -> newSTRef 0 + , accum = \bs i s -> + modifySTRef' s (+ (e (i,bs))) + , done = readSTRef + } +{-# INLINE r_add #-} + +r_nop :: Reducer xs s () +r_nop = Reducer + { init = \_ -> return () + , accum = \_ _ _ -> return () + , done = \_ -> return () + } +{-# INLINE r_nop #-} + pair :: a -> b -> (a, b) pair = (,) +{-# INLINE pair #-} true, false :: Bool true = True false = False +nothing :: Maybe a +nothing = Nothing + +just :: a -> Maybe a +just = Just + +left :: a -> Either a b +left = Left + +right :: b -> Either a b +right = Right + unit :: () unit = () @@ -104,36 +209,76 @@ ptrue, pfalse :: a -> Branch Bool a ptrue b = Branch { extract = extractBool True b } pfalse b = Branch { extract = extractBool False b } +{-# INLINE ptrue #-} +{-# INLINE pfalse #-} extractBool :: Bool -> a -> Bool -> Maybe a -extractBool b a p | p == b = Just a +extractBool b a p | p == b = Just a | otherwise = Nothing +{-# INLINE extractBool #-} + +pnothing :: b -> Branch (Maybe a) b +pnothing b = Branch { extract = \ma -> case ma of + Nothing -> Just b + Just _ -> Nothing } + +pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b +pjust PVar c = Branch { extract = \ma -> case ma of + Nothing -> Nothing + Just x -> Just (c x) } +pjust _ _ = error "TODO: Runtime.Prelude{pjust}" + +pleft :: Pattern -> (a -> c) -> Branch (Either a b) c +pleft PVar f = Branch { extract = \ma -> case ma of + Right _ -> Nothing + Left x -> Just (f x) } +pleft _ _ = error "TODO: Runtime.Prelude{pLeft}" + +pright :: Pattern -> (b -> c) -> Branch (Either a b) c +pright PVar f = Branch { extract = \ma -> case ma of + Left _ -> Nothing + Right x -> Just (f x) } +pright _ _ = error "TODO: Runtime.Prelude{pRight}" + + ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x,y) b ppair PVar PVar c = Branch { extract = (\(x,y) -> Just (c x y)) } ppair _ _ _ = error "ppair: TODO" +uncase_ :: Maybe a -> a +uncase_ (Just a) = a +uncase_ Nothing = error "case_: unable to match any branches" +{-# INLINE uncase_ #-} + case_ :: a -> [Branch a b] -> b -case_ e_ bs_ = go e_ bs_ - where go _ [] = error "case_: unable to match any branches" - go e (b:bs) = case extract b e of - Just b' -> b' - Nothing -> go e bs +case_ e [c1] = uncase_ (extract c1 e) +case_ e [c1, c2] = uncase_ (extract c1 e `mplus` extract c2 e) +case_ e bs_ = go bs_ + where go [] = error "case_: unable to match any branches" + go (b:bs) = case extract b e of + Just b' -> b' + Nothing -> go bs +{-# INLINE case_ #-} branch :: (c -> Branch a b) -> c -> Branch a b branch pat body = pat body +{-# INLINE branch #-} dirac :: a -> Measure a dirac = return +{-# INLINE dirac #-} pose :: Double -> Measure a -> Measure a pose _ a = a +{-# INLINE pose #-} superpose :: [(Double, Measure a)] -> Measure a superpose pms = do i <- makeMeasure $ MWCD.categorical (U.fromList $ map fst pms) snd (pms !! i) +{-# INLINE superpose #-} reject :: Measure a reject = Measure $ \_ -> return Nothing @@ -179,6 +324,7 @@ thRootOf :: Int -> Double -> Double thRootOf a b = b ** (recip $ fromIntegral a) +{-# INLINE thRootOf #-} array :: (G.Vector (MayBoxVec a) a) @@ -186,13 +332,29 @@ -> (Int -> a) -> MayBoxVec a a array n f = G.generate (fromIntegral n) (f . fromIntegral) +{-# INLINE array #-} +arrayLit :: (G.Vector (MayBoxVec a) a) => [a] -> MayBoxVec a a +arrayLit = G.fromList +{-# INLINE arrayLit #-} + (!) :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int -> a a ! b = a G.! (fromIntegral b) +{-# INLINE (!) #-} size :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int size v = fromIntegral (G.length v) +{-# INLINE size #-} +reduce + :: (G.Vector (MayBoxVec a) a) + => (a -> a -> a) + -> a + -> MayBoxVec a a + -> a +reduce f n v = G.foldr f n v +{-# INLINE reduce #-} + product :: Num a => Int @@ -200,6 +362,7 @@ -> (Int -> a) -> a product a b f = F.foldl' (\x y -> x * f y) 1 [a .. b-1] +{-# INLINE product #-} summate :: Num a @@ -208,6 +371,7 @@ -> (Int -> a) -> a summate a b f = F.foldl' (\x y -> x + f y) 0 [a .. b-1] +{-# INLINE summate #-} run :: Show a => MWC.GenIO diff --git a/haskell/Language/Hakaru/Sample.hs b/haskell/Language/Hakaru/Sample.hs --- a/haskell/Language/Hakaru/Sample.hs +++ b/haskell/Language/Hakaru/Sample.hs @@ -20,15 +20,20 @@ -- import qualified Numeric.Integration.TanhSinh as TS import qualified System.Random.MWC as MWC import qualified System.Random.MWC.Distributions as MWCD + import qualified Data.Vector as V +import Data.STRef import Data.Sequence (Seq) import qualified Data.Foldable as F import qualified Data.List.NonEmpty as L import Data.List.NonEmpty (NonEmpty(..)) import Data.Maybe (fromMaybe) + #if __GLASGOW_HASKELL__ < 710 import Control.Applicative (Applicative(..), (<$>)) #endif +import Control.Monad +import Control.Monad.ST import Control.Monad.Identity import Control.Monad.Trans.Maybe import Control.Monad.State.Strict @@ -43,6 +48,7 @@ import Language.Hakaru.Syntax.IClasses import Language.Hakaru.Syntax.TypeOf import Language.Hakaru.Syntax.Value +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.Datum import Language.Hakaru.Syntax.DatumCase import Language.Hakaru.Syntax.AST @@ -60,6 +66,15 @@ updateEnv v@(EAssoc x _) (Env xs) = Env $ IM.insert (fromNat $ varID x) v xs +updateEnvs + :: List1 Variable xs + -> List1 Value xs + -> Env + -> Env +updateEnvs Nil1 Nil1 env = env +updateEnvs (Cons1 x xs) (Cons1 y ys) env = + updateEnvs xs ys (updateEnv (EAssoc x y) env) + lookupVar :: Variable a -> Env -> Maybe (Value a) lookupVar x (Env env) = do EAssoc x' e' <- IM.lookup (fromNat $ varID x) env @@ -155,15 +170,17 @@ -> Value a evaluateTerm t env = case t of - o :$ es -> evaluateSCon o es env - NaryOp_ o es -> evaluateNaryOp o es env - Literal_ v -> evaluateLiteral v - Empty_ _ -> evaluateEmpty - Array_ n es -> evaluateArray n es env - Datum_ d -> evaluateDatum d env - Case_ o es -> evaluateCase o es env - Superpose_ es -> evaluateSuperpose es env - Reject_ _ -> VMeasure $ \_ _ -> return Nothing + o :$ es -> evaluateSCon o es env + NaryOp_ o es -> evaluateNaryOp o es env + Literal_ v -> evaluateLiteral v + Empty_ _ -> evaluateEmpty + Array_ n es -> evaluateArray n es env + ArrayLiteral_ es -> VArray . V.fromList $ map (flip evaluate env) es + Bucket b e rs -> evaluateBucket b e rs env + Datum_ d -> evaluateDatum d env + Case_ o es -> evaluateCase o es env + Superpose_ es -> evaluateSuperpose es env + Reject_ _ -> VMeasure $ \_ _ -> return Nothing evaluateSCon :: (ABT Term abt) @@ -536,6 +553,84 @@ VArray $ V.generate (fromNat n') $ \v -> let v' = VNat $ unsafeNat v in evaluate e' (updateEnv (EAssoc x v') env) + +evaluateBucket + :: (ABT Term abt) + => abt '[] 'HNat + -> abt '[] 'HNat + -> Reducer abt '[] a + -> Env + -> Value a +evaluateBucket b e rs env = + case (evaluate b env, evaluate e env) of + (VNat b', VNat e') -> runST $ do + s' <- init Nil1 rs env + mapM_ (\i -> accum (VNat i) Nil1 rs s' env) [b' .. e' - 1] + done s' + v2 -> case v2 of {} + where init :: (ABT Term abt) + => List1 Value xs + -> Reducer abt xs a + -> Env + -> ST s (VReducer s a) + init ix (Red_Fanout r1 r2) env = + VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env + init ix (Red_Index n _ mr) env = + let (vars, n') = caseBinds n in + case evaluate n' (updateEnvs vars ix env) of + VNat n'' -> VRed_Array <$> V.generateM (fromIntegral n'') + (\b -> init (Cons1 (vnat b) ix) mr env) + init ix (Red_Split _ r1 r2) env = + VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env + init ix Red_Nop env = return VRed_Unit + init ix (Red_Add h _) env = VRed_Num <$> newSTRef (identityElement (Sum h)) + + type_ = typeOfReducer + + vnat :: Int -> Value 'HNat + vnat = VNat . fromIntegral + + accum :: (ABT Term abt) + => Value 'HNat + -> List1 Value xs + -> Reducer abt xs a + -> VReducer s a + -> Env + -> ST s () + accum n ix (Red_Fanout r1 r2) (VRed_Pair s1 s2 v1 v2) env = + accum n ix r1 v1 env >> accum n ix r2 v2 env + accum n ix (Red_Index n' a1 r2) (VRed_Array v) env = + caseBind a1 $ \i a1' -> + let (vars, a1'') = caseBinds a1' + VNat ov = evaluate a1'' + (updateEnv (EAssoc i n) (updateEnvs vars ix env)) + ov' = fromIntegral ov in + accum n (Cons1 (VNat ov) ix) r2 (v V.! ov') env + accum n ix (Red_Split b r1 r2) (VRed_Pair s1 s2 v1 v2) env = + caseBind b $ \i b' -> + let (vars, b'') = caseBinds b' in + case evaluate b'' + (updateEnv (EAssoc i n) (updateEnvs vars ix env)) of + VDatum b' -> if b' == dTrue then + accum n ix r1 v1 env + else + accum n ix r2 v2 env + accum n ix (Red_Add h e) (VRed_Num s) env = + caseBind e $ \i e' -> + let (vars, e'') = caseBinds e' + v = evaluate e'' + (updateEnv (EAssoc i n) (updateEnvs vars ix env)) in + modifySTRef' s (evalOp (Sum h) v) + accum _ _ Red_Nop _ _ = return () + + done :: VReducer s a -> ST s (Value a) + done (VRed_Num s) = readSTRef s + done VRed_Unit = return (VDatum dUnit) + done (VRed_Pair s1 s2 v1 v2) = do + v1' <- done v1 + v2' <- done v2 + return (VDatum $ dPair_ s1 s2 v1' v2') + done (VRed_Array v) = VArray <$> V.sequence (V.map done v) evaluateDatum :: (ABT Term abt) diff --git a/haskell/Language/Hakaru/Simplify.hs b/haskell/Language/Hakaru/Simplify.hs --- a/haskell/Language/Hakaru/Simplify.hs +++ b/haskell/Language/Hakaru/Simplify.hs @@ -8,6 +8,7 @@ , OverloadedStrings , ScopedTypeVariables , TypeOperators + , RecordWildCards #-} {-# OPTIONS_GHC -Wall -fwarn-tabs #-} @@ -26,83 +27,29 @@ module Language.Hakaru.Simplify ( simplify , simplifyDebug - , MapleException(MapleException) ) where -import Control.Exception -import Control.Monad (when) - -import qualified Language.Hakaru.Pretty.Maple as Maple - -import Language.Hakaru.Parser.Maple -import Language.Hakaru.Parser.AST (Name) -import qualified Language.Hakaru.Parser.SymbolResolve as SR (resolveAST', fromVarSet) - import Language.Hakaru.Syntax.ABT import Language.Hakaru.Syntax.AST -import Language.Hakaru.Syntax.TypeCheck -import Language.Hakaru.Syntax.TypeOf - -import Language.Hakaru.Evaluation.ConstantPropagation - -import Data.Typeable (Typeable) - -import Data.Text (pack) -import System.MapleSSH (maple) -import System.IO +import Language.Hakaru.Syntax.Command +import Language.Hakaru.Maple ---------------------------------------------------------------- -data MapleException = MapleException String String - deriving Typeable - --- Maple prints errors with "cursors" (^) which point to the specific position --- of the error on the line above. The derived show instance doesn't preserve --- positioning of the cursor. -instance Show MapleException where - show (MapleException toMaple_ fromMaple) = - "MapleException:\n" ++ fromMaple ++ - "\nafter sending to Maple:\n" ++ toMaple_ - -instance Exception MapleException - simplify :: forall abt a . (ABT Term abt) => abt '[] a -> IO (abt '[] a) -simplify = simplifyDebug False +simplify = sendToMaple defaultMapleOptions{command=Simplify} simplifyDebug :: forall abt a . (ABT Term abt) - => Bool -> abt '[] a -> IO (abt '[] a) -simplifyDebug debug e = do - let typ = typeOf e - let toMaple_ = "use Hakaru, NewSLO in timelimit(90, RoundTrip(" - ++ Maple.pretty e ++ ", " ++ Maple.mapleType typ ")) end use;" - when debug (hPutStrLn stderr ("Sent to Maple:\n" ++ toMaple_)) - fromMaple <- maple toMaple_ - case fromMaple of - '_':'I':'n':'e':'r':'t':_ -> do - when debug $ do - ret <- maple ("FromInert(" ++ fromMaple ++ ")") - hPutStrLn stderr ("Returning from Maple:\n" ++ ret) - either (throw . MapleException toMaple_) - (return . constantPropagation) $ do - past <- leftShow $ parseMaple (pack fromMaple) - let m = checkType typ - (SR.resolveAST' (getNames e) (maple2AST past)) - leftShow $ unTCM m (freeVars e) Nothing UnsafeMode - _ -> throw (MapleException toMaple_ fromMaple) - - where - leftShow :: forall b c. Show b => Either b c -> Either String c - leftShow (Left err) = Left (show err) - leftShow (Right x) = Right x - - getNames :: abt '[] a -> [Name] - getNames = SR.fromVarSet . freeVars - + => Bool + -> Int + -> abt '[] a + -> IO (abt '[] a) +simplifyDebug d t = sendToMaple defaultMapleOptions{command=Simplify,debug=d,timelimit=t} ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/haskell/Language/Hakaru/Summary.hs b/haskell/Language/Hakaru/Summary.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Summary.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE TypeSynonymInstances + , FlexibleInstances + , FlexibleContexts + , DeriveDataTypeable + , CPP + , GADTs + , DataKinds + , OverloadedStrings + , ScopedTypeVariables + , TypeOperators + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.06.19 +-- | +-- Module : Language.Hakaru.Summary +-- Copyright : Copyright (c) 2017 the Hakaru team +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : experimental +-- Portability : GHC-only +-- +-- Take strings from Maple and interpret them in Haskell (Hakaru) +---------------------------------------------------------------- +module Language.Hakaru.Summary + ( summary + , summaryDebug + , MapleException(MapleInterpreterException) + ) where + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.Command +import Language.Hakaru.Maple + +---------------------------------------------------------------- + +summary + :: forall abt a + . (ABT Term abt) + => abt '[] a -> IO (abt '[] a) +summary = sendToMaple defaultMapleOptions{command=Summarize} + +summaryDebug + :: forall abt a + . (ABT Term abt) + => Bool -> abt '[] a -> IO (abt '[] a) +summaryDebug d = sendToMaple defaultMapleOptions{command=Summarize,debug=d} + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/haskell/Language/Hakaru/Syntax/ABT.hs b/haskell/Language/Hakaru/Syntax/ABT.hs --- a/haskell/Language/Hakaru/Syntax/ABT.hs +++ b/haskell/Language/Hakaru/Syntax/ABT.hs @@ -57,16 +57,20 @@ , rename , renames , subst + , substM , substs -- ** Constructing first-order trees with a HOAS-like API -- cf., , binder + , binderM + , Binders(binders) -- *** Highly experimental -- , Hint(..) -- , multibinder , withMetadata -- ** Abstract nonsense , cataABT + , cataABTM , paraABT -- * Some ABT instances @@ -75,13 +79,17 @@ , MetaABT(..) ) where -import Data.Text (Text) +import Data.Text (Text, empty) --import qualified Data.IntMap as IM import qualified Data.Foldable as F #if __GLASGOW_HASKELL__ < 710 -import Data.Monoid (Monoid(..)) +import Control.Applicative hiding (empty) +import Data.Monoid (Monoid(..)) #endif +import Control.Monad +import Control.Monad.Identity +import Control.Monad.Fix import Data.Number.Nat import Language.Hakaru.Syntax.IClasses -- TODO: factor the definition of the 'Sing' type family out from @@ -157,7 +165,17 @@ fmap12 _ (Var x) = Var x fmap12 f (Bind x e) = Bind x (fmap12 f e) +instance Foldable12 View where + foldMap12 f (Syn t) = f t + foldMap12 f (Var x) = mempty + foldMap12 f (Bind x e) = foldMap12 f e +instance Traversable12 View where + traverse12 f (Syn t) = Syn <$> f t + traverse12 _ (Var x) = pure $ Var x + traverse12 f (Bind x e) = Bind x <$> traverse12 f e + + instance (Show1 (Sing :: k -> *), Show1 rec) => Show2 (View (rec :: k -> *)) where @@ -739,49 +757,64 @@ -- should have strict 'fmap21' definitions. subst :: forall syn abt (a :: k) xs (b :: k) - . (JmEq1 (Sing :: k -> *), Show1 (Sing :: k -> *), Functor21 syn, ABT syn abt) + . (JmEq1 (Sing :: k -> *), Show1 (Sing :: k -> *), Traversable21 syn, ABT syn abt) => Variable a -> abt '[] a -> abt xs b -> abt xs b -subst x e = +subst x e = #ifdef __TRACE_DISINTEGRATE__ trace ("about to subst " ++ show (varID x)) $ -#endif +#endif + runIdentity . substM x e varCase + where + varCase :: forall m b'. (Applicative m, Functor m, Monad m) + => Variable b' -> m (abt '[] b') + varCase = return . var + +substM + :: forall syn abt (a :: k) xs (b :: k) m + . (JmEq1 (Sing :: k -> *), Show1 (Sing :: k -> *), + Traversable21 syn, ABT syn abt, + Applicative m, Functor m, Monad m) + => Variable a + -> abt '[] a + -> (forall b'. Variable b' -> m (abt '[] b')) + -> abt xs b + -> m (abt xs b) +substM x e vf = start (maxNextFreeOrBind [Some2 (var x), Some2 e]) where - -- TODO: we could use the director-strings approach to optimize this (for MemoizedABT, but pessimizing for TrivialABT) by first checking whether @x@ is free in @f@; if so then recurse, if not then we're done. - start :: forall xs' b'. Nat -> abt xs' b' -> abt xs' b' + -- TODO: we could use the director-strings approach to optimize + -- this (for MemoizedABT, but pessimizing for TrivialABT) by first + -- checking whether @x@ is free in @f@; if so then recurse, if not + -- then we're done. + start :: forall xs' b'. Nat -> abt xs' b' -> m (abt xs' b') start n f = loop n f (viewABT f) -- TODO: is it actually worth passing around the @f@? Benchmark. - loop :: forall xs' b'. Nat -> abt xs' b' -> View (syn abt) xs' b' -> abt xs' b' - loop n _ (Syn t) = syn $! fmap21 (start n) t + loop :: forall xs' b' + . Nat -> abt xs' b' -> View (syn abt) xs' b' -> m (abt xs' b') + loop n _ (Syn t) = syn <$> traverse21 (start n) t loop _ f (Var z) = #ifdef __TRACE_DISINTEGRATE__ trace ("checking varEq " ++ show (varID x) ++ " " ++ show (varID z)) $ #endif case varEq x z of - Just Refl -> e - Nothing -> f - loop n f (Bind z _) - | varID x == varID z = f - | otherwise = - -- TODO: even if we don't come up with a smarter way - -- of freshening variables, it'd be better to just pass - -- both sets to 'freshen' directly and then check them - -- each; rather than paying for taking their union every - -- time we go under a binder like this. - let i = 1 + max n (nextFreeOrBind f) -- (freeVars e `mappend` freeVars f) - z' = i `seq` z{varID = i} - -- HACK: the 'rename' function requires an ABT not a - -- View, so we have to use 'caseBind' to give its - -- input and then 'viewABT' to discard the topmost - -- annotation. We really should find a way to eliminate - -- that overhead. - in caseBind f $ \_ f' -> - let f'' = rename z z' f' in - bind z' (loop i f'' (viewABT f'')) + Just Refl -> return e + Nothing -> vf z + loop n f (Bind z b) = + if (varID x == varID z) then return f + else do -- TODO: even if we don't come up with a smarter way + -- of freshening variables, it'd be better to just pass + -- both sets to 'freshen' directly and then check them + -- each; rather than paying for taking their union every + -- time we go under a binder like this. + let i = 1 + max n (nextFreeOrBind f) -- (freeVars e `mappend` freeVars f) + z' = i `seq` z{varID = i} + f'' <- substM z (var z') vf (unviewABT b) + bind z' <$> loop i f'' (viewABT f'') + renames :: forall @@ -830,7 +863,7 @@ . ( ABT syn abt , JmEq1 (Sing :: k -> *) , Show1 (Sing :: k -> *) - , Functor21 syn + , Traversable21 syn ) => Assocs (abt '[]) -> abt xs a @@ -915,6 +948,36 @@ x = Variable hint (nextBind body) typ -- N.B., cannot use 'nextFree' when deciding the 'varID' of @x@ + +-- A Monadic variant of @binder@ which allows constructing a term in a monadic +-- context. The dependency on MonadFix is due to the knot-tying used to generate +-- the bound variable. +binderM + :: (MonadFix m, ABT syn abt) + => Text + -> Sing a + -> (abt '[] a -> m (abt xs b)) + -> m (abt (a ': xs) b) +binderM hint typ hoas = do + (var, body) <- mfix $ \ ~(_, b) -> do + let v = Variable hint (nextBind b) typ + b' <- hoas (var v) + return (v, b') + return (bind var body) + +class (ABT syn abt) => + Binders syn abt xs as | abt -> syn, abt xs -> as, abt as -> xs where + binders :: (as -> abt '[] b) -> abt xs b + +instance (ABT syn abt) => + Binders syn abt '[] () where + binders hoas = hoas () + +instance (Binders syn abt xs as, SingI x) => + Binders syn abt (x ': xs) (abt '[] x, as) where + binders hoas = binder empty sing (binders . curry hoas) + + {- data Hint (a :: k) = Hint !Text !(Sing a) @@ -986,6 +1049,31 @@ loop (Bind x e) = bind_ x (loop e) {-# INLINE cataABT #-} +-- | A monadic variant of @cataABT@, which may not fit the precise definition of +-- a catamorphism? The @bind@ and @syn@ operations receive monadic actions as their +-- inputs, which allows the @bind@ and @syn@ operations to update the monadic +-- context in which the subterms are evaluated if need be. +cataABTM + :: forall + (abt :: [k] -> k -> *) + (syn :: ([k] -> k -> *) -> k -> *) + (r :: [k] -> k -> *) + (f :: * -> *) + . (ABT syn abt, Traversable21 syn, Applicative f) + => (forall a. Variable a -> f (r '[] a)) + -> (forall x xs a. Variable x -> f (r xs a) -> f (r (x ': xs) a)) + -> (forall a. f (syn r a) -> f (r '[] a)) + -> forall xs a. abt xs a -> f (r xs a) +cataABTM var_ bind_ syn_ = start + where + start :: forall ys b. abt ys b -> f (r ys b) + start = loop . viewABT + + loop :: forall ys b. View (syn abt) ys b -> f (r ys b) + loop (Syn t) = syn_ (traverse21 start t) + loop (Var x) = var_ x + loop (Bind x e) = bind_ x (loop e) +{-# INLINE cataABTM #-} -- | The paramorphism (aka: recursor) for ABTs. While this is -- equivalent to 'cataABT' in terms of the definable /functions/, diff --git a/haskell/Language/Hakaru/Syntax/ANF.hs b/haskell/Language/Hakaru/Syntax/ANF.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/ANF.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.ANF +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.ANF (normalize, isValue) where + +import Data.Maybe +import Control.Monad.Cont (Cont, runCont, cont) + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Types.DataKind + +import Language.Hakaru.Syntax.Prelude + +-- The renaming environment which maps variables in the original term to their +-- counterparts in the new term. This is needed since the mechanism which +-- ensures hygiene for the AST only factors in binders, but not free variables +-- in the expression being constructed. When we construct a new binding form, a +-- new variable is introduced and the variable in the old expression must be +-- mapped to the new one. + +type Env = Assocs (Variable :: Hakaru -> *) + +updateEnv :: forall (a :: Hakaru) . Variable a -> Variable a -> Env -> Env +updateEnv vin vout = insertAssoc (Assoc vin vout) + +-- | The context in which A-normalization occurs. Represented as a continuation, +-- the context expects an expression of a particular type (usually a variable) +-- and produces a new expression as a result. +type Context abt a b = abt '[] a -> abt '[] b + +-- | Extract a variable from an abt. This function is partial +getVar :: (ABT Term abt) => abt xs a -> Variable a +getVar abt = case viewABT abt of + Var v -> v + _ -> error "getVar: not given a variable" + +-- | Useful function for generating fresh variables from an existing variable by +-- wrapping @binder@. +freshVar + :: (ABT Term abt) + => Variable a + -> (Variable a -> abt xs b) + -> abt (a ': xs) b +freshVar v f = binder (varHint v) (varType v) (f . getVar) + +remapVar + :: (ABT Term abt) + => Variable a + -> Env + -> (Env -> abt xs b) + -> abt (a ': xs) b +remapVar v env f = freshVar v $ \v' -> f (updateEnv v v' env) + +-- | Entry point for the normalization process. Initializes normalize' with the +-- empty context. +normalize + :: (ABT Term abt) + => abt '[] a + -> abt '[] a +normalize abt = normalize' abt emptyAssocs id + +normalize' + :: (ABT Term abt) + => abt '[] a + -> Env + -> Context abt a b + -> abt '[] b +normalize' = normalizeTail . viewABT + +normalizeTail, normalizeSave + :: (ABT Term abt) + => View (Term abt) xs a + -> Env + -> (abt xs a -> abt '[] b) + -> abt '[] b +normalizeTail (Var v) env ctxt = ctxt (normalizeVar v env) +normalizeTail (Syn s) env ctxt = normalizeTerm s env ctxt +normalizeTail view@Bind{} env ctxt = ctxt (normalizeReset view env) +normalizeSave (Var v) env ctxt = ctxt (normalizeVar v env) +normalizeSave (Syn s) env ctxt = normalizeTerm s env giveName + where giveName abt' | isValue abt' = ctxt abt' + | otherwise = let_ abt' ctxt +normalizeSave view@Bind{} env ctxt = ctxt (normalizeReset view env) + +normalizeReset :: (ABT Term abt) => View (Term abt) xs a -> Env -> abt xs a +normalizeReset (Var v) env = normalizeVar v env +normalizeReset (Syn s) env = normalizeTerm s env id +normalizeReset (Bind v b) env = remapVar v env (normalizeReset b) + +normalizeVar :: (ABT Term abt) => Variable a -> Env -> abt '[] a +normalizeVar v env = var $ fromMaybe v (lookupAssoc v env) + +isValue + :: (ABT Term abt) + => abt xs a + -> Bool +isValue abt = + case viewABT abt of + Var{} -> True + Bind{} -> False + Syn s -> isValueTerm s + where + isValueTerm Literal_{} = True + isValueTerm (Lam_ :$ _) = True + isValueTerm _ = False + +normalizeTerm + :: forall abt a b + . (ABT Term abt) + => Term abt a + -> Env + -> Context abt a b + -> abt '[] b + +normalizeTerm (Let_ :$ (rhs :* body :* End)) env ctxt = + caseBind body $ \v body' -> + normalize' rhs env $ \rhs' -> + let mkbody env' = normalize' body' env' ctxt + in syn (Let_ :$ rhs' :* remapVar v env mkbody :* End) + +normalizeTerm (Case_ cond bs) env ctxt = + normalizeSave (viewABT cond) env $ \ cond' -> + let normalizeBranch :: forall xs d . abt xs d -> abt xs d + normalizeBranch body = normalizeReset (viewABT body) env + branches = map (fmap21 normalizeBranch) bs + -- A possible optimization is to push the context into each conditional, + -- possibly opening up other optimizations at the cost of code growth. + in ctxt $ syn $ Case_ cond' branches + +normalizeTerm term env ctxt = runCont (fmap syn (traverse21 f term)) ctxt + where f :: forall xs c . abt xs c -> Cont (abt '[] b) (abt xs c) + f abt = cont (n (viewABT abt) env) + n :: forall xs c + . View (Term abt) xs c + -> Env + -> (abt xs c -> abt '[] b) + -> abt '[] b + -- TODO: Can we just let n=normalizeTail or n=normalizeSave? + n = case term of MBind :$ _ -> normalizeTail + Plate :$ _ -> normalizeTail + Dirac :$ _ -> normalizeTail + UnsafeFrom_ _ :$ _ -> normalizeTail + CoerceTo_ _ :$ _ -> normalizeTail + _ -> normalizeSave diff --git a/haskell/Language/Hakaru/Syntax/AST.hs b/haskell/Language/Hakaru/Syntax/AST.hs --- a/haskell/Language/Hakaru/Syntax/AST.hs +++ b/haskell/Language/Hakaru/Syntax/AST.hs @@ -74,6 +74,7 @@ import Language.Hakaru.Types.HClasses import Language.Hakaru.Types.Coercion import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.ABT (ABT(syn)) ---------------------------------------------------------------- @@ -163,22 +164,22 @@ int2nat x = case toNatural x of Just y -> y - Nothing -> error "primCoerceFrom@Literal: negative HInt" + Nothing -> error $ "primCoerceFrom@Literal: negative HInt " ++ show x prob2nat :: NonNegativeRational -> Natural prob2nat x = if denominator x == 1 then numerator x - else error "primCoerceFrom@Literal: non-integral HProb" + else error $ "primCoerceFrom@Literal: non-integral HProb " ++ show x real2prob :: Rational -> NonNegativeRational real2prob x = case toNonNegativeRational x of Just y -> y - Nothing -> error "primCoerceFrom@Literal: negative HReal" + Nothing -> error $ "primCoerceFrom@Literal: negative HReal " ++ show x real2int :: Rational -> Integer real2int x = if denominator x == 1 then numerator x - else error "primCoerceFrom@Literal: non-integral HReal" + else error $ "primCoerceFrom@Literal: non-integral HReal " ++ show x ---------------------------------------------------------------- @@ -784,6 +785,17 @@ -> !(abt '[ 'HNat ] a) -> Term abt ('HArray a) + ArrayLiteral_ + :: [abt '[] a] + -> Term abt ('HArray a) + + -- Constructor for Reducers + Bucket + :: !(abt '[] 'HNat) + -> !(abt '[] 'HNat) + -> Reducer abt '[] a + -> Term abt a + -- -- User-defined data types -- BUG: even though the 'Datum' type has a single constructor, we get a warning about not being able to UNPACK it in 'Datum_'... wtf? -- @@ -801,6 +813,7 @@ :: L.NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a)) -> Term abt ('HMeasure a) + -- The zero measure Reject_ :: !(Sing ('HMeasure a)) -> Term abt ('HMeasure a) ---------------------------------------------------------------- @@ -857,17 +870,19 @@ . showList2 (F.toList es) ) ) - Literal_ v -> showParen_0 p "Literal_" v - Empty_ _ -> showString "Empty_" - Array_ e1 e2 -> showParen_22 p "Array_" e1 e2 - Datum_ d -> showParen_1 p "Datum_" (fmap11 LC_ d) - Case_ e bs -> + Literal_ v -> showParen_0 p "Literal_" v + Empty_ _ -> showString "Empty_" + Array_ e1 e2 -> showParen_22 p "Array_" e1 e2 + ArrayLiteral_ es -> showParen (p > 9) (showString "ArrayLiteral_" . showList2 es) + Datum_ d -> showParen_1 p "Datum_" (fmap11 LC_ d) + Case_ e bs -> showParen (p > 9) ( showString "Case_ " . showsPrec2 11 e . showString " " . showList1 bs ) + Bucket _ _ _ -> showString "Bucket ..." Superpose_ pes -> showParen (p > 9) ( showString "Superpose_ " @@ -889,8 +904,10 @@ fmap21 _ (Literal_ v) = Literal_ v fmap21 _ (Empty_ t) = Empty_ t fmap21 f (Array_ e1 e2) = Array_ (f e1) (f e2) + fmap21 f (ArrayLiteral_ es) = ArrayLiteral_ (fmap f es) fmap21 f (Datum_ d) = Datum_ (fmap11 f d) fmap21 f (Case_ e bs) = Case_ (f e) (map (fmap21 f) bs) + fmap21 f (Bucket b e r) = Bucket (f b) (f e) (fmap22 f r) fmap21 f (Superpose_ pes) = Superpose_ (L.map (f *** f) pes) fmap21 _ (Reject_ t) = Reject_ t @@ -902,8 +919,10 @@ foldMap21 _ (Literal_ _) = mempty foldMap21 _ (Empty_ _) = mempty foldMap21 f (Array_ e1 e2) = f e1 `mappend` f e2 + foldMap21 f (ArrayLiteral_ es) = F.foldMap f es foldMap21 f (Datum_ d) = foldMap11 f d - foldMap21 f (Case_ e bs) = f e `mappend` F.foldMap (foldMap21 f) bs + foldMap21 f (Case_ e bs) = f e `mappend` F.foldMap (foldMap21 f) bs + foldMap21 f (Bucket b e r) = f b `mappend` f e `mappend` foldMap22 f r foldMap21 f (Superpose_ pes) = foldMapPairs f pes foldMap21 _ (Reject_ _) = mempty @@ -922,6 +941,8 @@ traverse21 _ (Literal_ v) = pure $ Literal_ v traverse21 _ (Empty_ typ) = pure $ Empty_ typ traverse21 f (Array_ e1 e2) = Array_ <$> f e1 <*> f e2 + traverse21 f (ArrayLiteral_ es) = ArrayLiteral_ <$> traverse f es + traverse21 f (Bucket b e r) = Bucket <$> f b <*> f e <*> traverse22 f r traverse21 f (Datum_ d) = Datum_ <$> traverse11 f d traverse21 f (Case_ e bs) = Case_ <$> f e <*> traverse (traverse21 f) bs traverse21 f (Superpose_ pes) = Superpose_ <$> traversePairs f pes diff --git a/haskell/Language/Hakaru/Syntax/AST/Eq.hs b/haskell/Language/Hakaru/Syntax/AST/Eq.hs --- a/haskell/Language/Hakaru/Syntax/AST/Eq.hs +++ b/haskell/Language/Hakaru/Syntax/AST/Eq.hs @@ -48,6 +48,7 @@ import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.Datum import Language.Hakaru.Syntax.TypeOf +import Language.Hakaru.Syntax.Reducer import Control.Monad.Reader @@ -63,6 +64,7 @@ import Data.Maybe +import Data.List (inits,tails) -- import Data.Number.Nat import Unsafe.Coerce @@ -121,11 +123,20 @@ Refl <- jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') Refl <- jmEq1 es es' Just (Refl, Refl) +jmEq_S (Product h1 h2) es (Product h1' h2') es' = do + Refl <- jmEq1 (sing_HDiscrete h1) (sing_HDiscrete h1') + Refl <- jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') + Refl <- jmEq1 es es' + Just (Refl, Refl) +jmEq_S (Product h1 h2) es (Product h1' h2') es' = do + Refl <- jmEq1 (sing_HDiscrete h1) (sing_HDiscrete h1') + Refl <- jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') + Refl <- jmEq1 es es' + Just (Refl, Refl) jmEq_S Expect es Expect es' = jmEq1 es es' >>= \Refl -> Just (Refl, Refl) jmEq_S _ _ _ _ = Nothing - -- TODO: Handle jmEq2 of pat and pat' jmEq_Branch :: (ABT Term abt, JmEq2 abt) @@ -162,6 +173,17 @@ (Refl, Refl) <- jmEq2 i j (Refl, Refl) <- jmEq2 f g Just Refl + -- Assumes nonempty literal arrays. The hope is that Empty_ covers that case. + -- TODO handle empty literal arrays. + jmEq1 (ArrayLiteral_ (e:es)) (ArrayLiteral_ (e':es')) = do + (Refl, Refl) <- jmEq2 e e' + () <- all_jmEq2 (S.fromList es) (S.fromList es') + return Refl + jmEq1 (Bucket a b r) (Bucket a' b' r') = do + (Refl, Refl) <- jmEq2 a a' + (Refl, Refl) <- jmEq2 b b' + Refl <- jmEq1 r r' + return Refl jmEq1 (Datum_ (Datum hint _ _)) (Datum_ (Datum hint' _ _)) -- BUG: We need to compare structurally rather than using the hint | hint == hint' = unsafeCoerce (Just Refl) @@ -273,6 +295,17 @@ try_bool :: Bool -> ReaderT Varmap Maybe () try_bool b = lift $ if b then Just () else Nothing +split :: [a] -> [(a,[a])] +split xs = zipWith (\as (b:bs)->(b,as++bs)) (inits xs) (init $ tails xs) + +zipWithSetM :: MonadPlus m => (a -> a -> m ()) -> [a] -> [a] -> m () +zipWithSetM _ [] ys = guard (null ys) +zipWithSetM q (x:xs) ys = msum [ q x y >> zipWithSetM q xs ys' | (y,ys') <- split ys ] + +zipWithSetMF :: (MonadPlus m, F.Foldable f) => (a -> a -> m ()) -> f a -> f a -> m () +zipWithSetMF q a b = zipWithSetM q (F.toList a) (F.toList b) + + alphaEq :: forall abt a . (ABT Term abt) @@ -316,19 +349,25 @@ (o1 :$ es1, o2 :$ es2) -> sConEq o1 es1 o2 es2 (NaryOp_ op1 es1, NaryOp_ op2 es2) -> do try_bool (op1 == op2) - F.sequence_ $ S.zipWith go (viewABT <$> es1) (viewABT <$> es2) + zipWithSetMF go (viewABT <$> es1) (viewABT <$> es2) (Literal_ x, Literal_ y) -> try_bool (x == y) (Empty_ x, Empty_ y) -> void_jmEq1 x y (Datum_ d1, Datum_ d2) -> datumEq d1 d2 (Array_ n1 e1, Array_ n2 e2) -> do go (viewABT n1) (viewABT n2) go (viewABT e1) (viewABT e2) + (ArrayLiteral_ es, ArrayLiteral_ es') -> + F.sequence_ $ zipWith go (viewABT <$> es) (viewABT <$> es') + (Bucket a b r, Bucket a' b' r') -> do + go (viewABT a) (viewABT a') + go (viewABT b) (viewABT b') + reducerEq r r' (Case_ e1 bs1, Case_ e2 bs2) -> do Refl <- lift $ jmEq1 (typeOf e1) (typeOf e2) go (viewABT e1) (viewABT e2) zipWithM_ sBranch bs1 bs2 (Superpose_ pms1, Superpose_ pms2) -> - F.sequence_ $ L.zipWith pairEq pms1 pms2 + zipWithSetMF pairEq pms1 pms2 (Reject_ x, Reject_ y) -> void_jmEq1 x y (_, _) -> lift Nothing @@ -366,12 +405,14 @@ go (viewABT e2) (viewABT e2') sConEq (CoerceTo_ _) (e1 :* End) - (CoerceTo_ _) (e2 :* End) = - void_jmEq1 (typeOf e1) (typeOf e2) + (CoerceTo_ _) (e2 :* End) = do + Refl <- lift $ jmEq1 (typeOf e1) (typeOf e2) + go (viewABT e1) (viewABT e2) sConEq (UnsafeFrom_ _) (e1 :* End) - (UnsafeFrom_ _) (e2 :* End) = - void_jmEq1 (typeOf e1) (typeOf e2) + (UnsafeFrom_ _) (e2 :* End) = do + Refl <- lift $ jmEq1 (typeOf e1) (typeOf e2) + go (viewABT e1) (viewABT e2) sConEq (PrimOp_ o1) es1 (PrimOp_ o2) es2 = primOpEq o1 es1 o2 es2 @@ -400,6 +441,16 @@ Refl <- lift $ jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') sArgsEq e1 e2 + sConEq (Product h1 h2) e1 (Product h1' h2') e2 = do + Refl <- lift $ jmEq1 (sing_HDiscrete h1) (sing_HDiscrete h1') + Refl <- lift $ jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') + sArgsEq e1 e2 + + sConEq (Product h1 h2) e1 (Product h1' h2') e2 = do + Refl <- lift $ jmEq1 (sing_HDiscrete h1) (sing_HDiscrete h1') + Refl <- lift $ jmEq1 (sing_HSemiring h2) (sing_HSemiring h2') + sArgsEq e1 e2 + sConEq Expect (e1 :* e2 :* End) Expect (e1' :* e2' :* End) = do Refl <- lift $ jmEq1 (typeOf e1) (typeOf e1') @@ -491,4 +542,59 @@ . Branch a abt b -> Branch a abt b -> ReaderT Varmap Maybe () - sBranch (Branch _ e1) (Branch _ e2) = go (viewABT e1) (viewABT e2) + sBranch (Branch p1 e1) (Branch p2 e2) = patternEq p1 p2 >> go (viewABT e1) (viewABT e2) + + patternEq + :: Pattern a0 b0 + -> Pattern a1 b1 + -> ReaderT Varmap Maybe () + patternEq PWild PWild = return () + patternEq PVar PVar = return () + patternEq (PDatum _ a) (PDatum _ b) = pdatumCodeEq a b + patternEq _ _ = mzero + + pdatumCodeEq + :: PDatumCode xss0 vs0 a0 + -> PDatumCode xss1 vs1 a1 + -> ReaderT Varmap Maybe () + pdatumCodeEq (PInr c) (PInr d) = pdatumCodeEq c d + pdatumCodeEq (PInl c) (PInl d) = pdatumStructEq c d + pdatumCodeEq _ _ = mzero + + pdatumStructEq + :: PDatumStruct xs0 vs0 a0 + -> PDatumStruct xs1 vs1 a1 + -> ReaderT Varmap Maybe () + pdatumStructEq (PEt c1 c2) (PEt d1 d2) = do + pdatumFunEq c1 d1 + pdatumStructEq c2 d2 + pdatumStructEq PDone PDone = return () + pdatumStructEq _ _ = lift Nothing + + pdatumFunEq + :: PDatumFun x0 vs0 a0 + -> PDatumFun x1 vs1 a1 + -> ReaderT Varmap Maybe () + pdatumFunEq (PKonst e) (PKonst f) = patternEq e f + pdatumFunEq (PIdent e) (PIdent f) = patternEq e f + pdatumFunEq _ _ = lift Nothing + + reducerEq + :: forall xs a + . Reducer abt xs a + -> Reducer abt xs a + -> ReaderT Varmap Maybe () + reducerEq (Red_Fanout r s) (Red_Fanout r' s') = do + reducerEq r r' + reducerEq s s' + reducerEq (Red_Index s i r) (Red_Index s' i' r') = do + go (viewABT s) (viewABT s') + go (viewABT i) (viewABT i') + reducerEq r r' + reducerEq (Red_Split i r s) (Red_Split i' r' s') = do + go (viewABT i) (viewABT i') + reducerEq r r' + reducerEq s s' + reducerEq Red_Nop Red_Nop = return () + reducerEq (Red_Add _ x) (Red_Add _ x') = go (viewABT x) (viewABT x') + reducerEq _ _ = lift Nothing diff --git a/haskell/Language/Hakaru/Syntax/AST/Transforms.hs b/haskell/Language/Hakaru/Syntax/AST/Transforms.hs --- a/haskell/Language/Hakaru/Syntax/AST/Transforms.hs +++ b/haskell/Language/Hakaru/Syntax/AST/Transforms.hs @@ -12,6 +12,11 @@ import qualified Data.Sequence as S +import Language.Hakaru.Syntax.ANF (normalize) +import Language.Hakaru.Syntax.CSE (cse) +import Language.Hakaru.Syntax.Prune (prune) +import Language.Hakaru.Syntax.Uniquify (uniquify) +import Language.Hakaru.Syntax.Hoist (hoist) import Language.Hakaru.Syntax.ABT import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.TypeOf @@ -20,6 +25,18 @@ import Language.Hakaru.Expect (expect) import Language.Hakaru.Disintegrate (determine, observe) + +optimizations + :: (ABT Term abt) + => abt '[] a + -> abt '[] a +optimizations = uniquify + . prune + . cse + . hoist + -- The hoist pass needs globally uniqiue identifiers + . uniquify + . normalize underLam :: (ABT Term abt, Monad m) diff --git a/haskell/Language/Hakaru/Syntax/CSE.hs b/haskell/Language/Hakaru/Syntax/CSE.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/CSE.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE CPP + , DataKinds + , EmptyCase + , ExistentialQuantification + , FlexibleContexts + , GADTs + , GeneralizedNewtypeDeriving + , KindSignatures + , MultiParamTypeClasses + , OverloadedStrings + , PolyKinds + , ScopedTypeVariables + , TypeFamilies + , TypeOperators + #-} + + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.CSE +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.CSE (cse) where + +import Control.Monad.Reader + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Syntax.TypeOf +import Language.Hakaru.Types.DataKind + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +-- What we need is an environment like data structure which maps Terms (or +-- general abts?) to other abts. Can such a mapping be implemented efficiently? +-- This would seem to require a hash operation to make efficient. + +data EAssoc (abt :: [Hakaru] -> Hakaru -> *) + = forall a . EAssoc !(abt '[] a) !(abt '[] a) + +-- An association list for now +newtype Env (abt :: [Hakaru] -> Hakaru -> *) = Env [EAssoc abt] + +emptyEnv :: Env a +emptyEnv = Env [] + +trivial :: (ABT Term abt) => abt '[] a -> Bool +trivial abt = case viewABT abt of + Var _ -> True + Syn (Literal_ _) -> True + _ -> False + +-- Attempt to find a new expression in the environment. The lookup is chained +-- to iteratively perform lookup until no match is found, resulting in an +-- equivalence-relation in the environment. This could be made faster with path +-- compression and a more efficient lookup structure. +-- NB: This code could potentially produce an infinite loop depending on how +-- terms are added to the environment. How do we want to prevent this? +lookupEnv + :: forall abt a . (ABT Term abt) + => abt '[] a + -> Env abt + -> abt '[] a +lookupEnv start (Env env) = go start env + where + go :: abt '[] a -> [EAssoc abt] -> abt '[] a + go ast [] = ast + go ast (EAssoc a b : xs) = + case jmEq1 (typeOf ast) (typeOf a) of + Just Refl | alphaEq ast a -> go b env + _ -> go ast xs + +insertEnv + :: forall abt a . (ABT Term abt) + => abt '[] a + -> abt '[] a + -> Env abt + -> Env abt +insertEnv ast1 ast2 (Env env) + -- Point new variables to the older ones, this does not affect the amount of + -- work done, since ast2 is always a variable. This allows the pass to + -- eliminate redundant variables, as we only eliminate binders during CSE. + | trivial ast1 = Env (EAssoc ast2 ast1 : env) + -- Otherwise map expressions to their binding variables + | otherwise = Env (EAssoc ast1 ast2 : env) + +newtype CSE (abt :: [Hakaru] -> Hakaru -> *) a = CSE { runCSE :: Reader (Env abt) a } + deriving (Functor, Applicative, Monad, MonadReader (Env abt)) + +replaceCSE + :: (ABT Term abt) + => abt '[] a + -> CSE abt (abt '[] a) +replaceCSE abt = lookupEnv abt `fmap` ask + +cse :: forall abt a . (ABT Term abt) => abt '[] a -> abt '[] a +cse abt = runReader (runCSE (cse' abt)) emptyEnv + +cse' :: forall abt xs a . (ABT Term abt) => abt xs a -> CSE abt (abt xs a) +cse' = loop . viewABT + where + loop :: View (Term abt) ys a -> CSE abt (abt ys a) + loop (Var v) = cseVar v + loop (Syn s) = cseTerm s + loop (Bind v b) = fmap (bind v) (loop b) + +-- Variables can be equivalent to other variables +-- TODO: A good sanity check would be to ensure the result in this case is +-- always a variable or constant. A variable should never be substituted for +-- a more complex expression. +cseVar + :: (ABT Term abt) + => Variable a + -> CSE abt (abt '[] a) +cseVar = replaceCSE . var + +mklet :: ABT Term abt => Variable b -> abt '[] b -> abt '[] a -> abt '[] a +mklet v rhs body = syn (Let_ :$ rhs :* bind v body :* End) + +-- Thanks to A-normalization, the only case we need to care about is let bindings. +-- Everything else is just structural recursion. +cseTerm + :: (ABT Term abt) + => Term abt a + -> CSE abt (abt '[] a) + +cseTerm (Let_ :$ rhs :* body :* End) = do + rhs' <- cse' rhs + caseBind body $ \v body' -> + local (insertEnv rhs' (var v)) $ + if trivial rhs' + then cse' body' + else fmap (mklet v rhs') (cse' body') + +cseTerm term = traverse21 cse' term >>= replaceCSE . syn + diff --git a/haskell/Language/Hakaru/Syntax/Command.hs b/haskell/Language/Hakaru/Syntax/Command.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Command.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE FlexibleInstances + , GADTs + , DataKinds + , TypeOperators + , ViewPatterns + , KindSignatures + , RankNTypes + , UndecidableInstances + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- | +-- Module : Language.Hakaru.Syntax.Command +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Stability : experimental +-- Portability : GHC-only +-- +-- An encoding of (some) Hakaru commands and their types. +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Command where + +import Language.Hakaru.Types.Sing +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.IClasses +import GHC.TypeLits (Symbol) +import Data.List (isInfixOf) +import Data.Char (toLower) +import Data.Function (on) + +---------------------------------------------------------------- + +data CommandType (c :: Symbol) (i :: Hakaru) (o :: Hakaru) where + Simplify :: CommandType "Simplify" a a + DisintMeas :: CommandType "Disintegrate" (HMeasure (HPair a b)) (a :-> HMeasure b) + DisintFun :: !(CommandType "Disintegrate" x x') + -> CommandType "Disintegrate" (a :-> x) (a :-> x') + Summarize :: CommandType "Summarize" a a + +commandIsType :: CommandType c i o -> Sing i -> Sing o +commandIsType DisintMeas (SMeasure (sUnPair->(a,b))) = SFun a (SMeasure b) +commandIsType (DisintFun t) (SFun a x) = SFun a (commandIsType t x) +commandIsType Simplify x = x +commandIsType Summarize x = x + +nameOfCommand :: CommandType c i o -> Sing c +nameOfCommand Simplify{} = sing +nameOfCommand Summarize{} = sing +nameOfCommand DisintMeas{} = sing +nameOfCommand DisintFun{} = sing + +parseCommand = flip (isInfixOf `on` map toLower) + +commandFromName + :: String + -> Sing i + -> (forall o c . Either Bool (CommandType c i o, Sing o) -> k) + -> k +commandFromName (parseCommand "Simplify"->True) i k = k $ Right (Simplify, i) + +commandFromName (parseCommand "Disintegrate"->True) i k = + let disint_commandFromType + :: Sing i + -> (forall o . Either Bool (CommandType "Disintegrate" i o, Sing o) -> k) + -> k + disint_commandFromType i k = + case i of + SMeasure (SData (STyApp (STyApp (STyCon (jmEq1 sSymbol_Pair -> Just Refl)) a) b) _) -> + k $ Right (DisintMeas, SFun a (SMeasure b)) + SFun a x -> + disint_commandFromType x $ \q -> + k $ fmap (\(c,x') -> (DisintFun c, SFun a x')) q + _ -> k $ Left True + in disint_commandFromType i k + +commandFromName (parseCommand "Summarize"->True) i k = k $ Right (Summarize, i) + +commandFromName _ _ k = k $ Left False diff --git a/haskell/Language/Hakaru/Syntax/Gensym.hs b/haskell/Language/Hakaru/Syntax/Gensym.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Gensym.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE CPP + , DataKinds + , FlexibleInstances + , FlexibleContexts + , KindSignatures + , OverloadedStrings + , UndecidableInstances + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Language.Hakaru.Syntax.Gensym where + +#if __GLASGOW_HASKELL__ < 710 +import Data.Functor ((<$>)) +#endif +import Control.Monad.State +import Data.Number.Nat +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.TypeOf +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.Sing + +class Monad m => Gensym m where + freshVarId :: m Nat + +instance (Monad m, MonadState Nat m) => Gensym m where + freshVarId = do + vid <- gets succ + put vid + return vid + +freshVar + :: (Functor m, Gensym m) + => Variable (a :: Hakaru) -> m (Variable a) +freshVar v = fmap (\n -> v{varID=n}) freshVarId + +varOfType + :: (Functor m, Gensym m) + => Sing (a :: Hakaru) -> m (Variable a) +varOfType t = fmap (\n -> Variable "" n t) freshVarId + +varForExpr + :: (Functor m, Gensym m, ABT Term abt) + => abt '[] a -> m (Variable a) +varForExpr v = fmap (\n -> Variable "" n (typeOf v)) freshVarId + diff --git a/haskell/Language/Hakaru/Syntax/Hoist.hs b/haskell/Language/Hakaru/Syntax/Hoist.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Hoist.hs @@ -0,0 +1,403 @@ +{-# LANGUAGE CPP + , BangPatterns + , DataKinds + , EmptyCase + , ExistentialQuantification + , FlexibleContexts + , FlexibleInstances + , GADTs + , GeneralizedNewtypeDeriving + , KindSignatures + , MultiParamTypeClasses + , OverloadedStrings + , PolyKinds + , ScopedTypeVariables + , StandaloneDeriving + , TupleSections + , TypeFamilies + , TypeOperators + , UndecidableInstances + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.Hoist +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- Hoist expressions to the point where their data dependencies are met. +-- This pass duplicates *a lot* of work and relies on a the CSE and pruning +-- passes to cleanup the junk (most of which is trivial to do, but we don't know +-- what is junk until after CSE has occured). +-- +-- NOTE: This pass assumes globally unique variable ids, as two subterms may +-- otherwise bind the same variable. Those variables would potentially shadow +-- eachother if hoisted upward to a common scope. +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Hoist (hoist) where + +import Control.Applicative (liftA2) +import Control.Monad.RWS +import qualified Data.Foldable as F +import qualified Data.Graph as G +import qualified Data.IntMap.Strict as IM +import qualified Data.List as L +import Data.Maybe (mapMaybe) +import Data.Number.Nat +import Data.Proxy (KProxy (..)) +import qualified Data.Vector as V + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.ANF (isValue) +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq (alphaEq) +import Language.Hakaru.Syntax.Gensym +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.Sing (Sing) + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +data Entry (abt :: Hakaru -> *) + = forall (a :: Hakaru) . Entry + { varDependencies :: !(VarSet (KindOf a)) + , expression :: !(abt a) + -- The type of the expression, to allow for easy comparison of types. + -- The typeOf operator is technically O(n) in the size of the expresion + -- and we may need to call it many times. + , sing :: !(Sing a) + , bindings :: ![Variable a] + } + +instance Show (Entry abt) where + show (Entry d _ _ b) = "Entry (" ++ show d ++ ") (" ++ show b ++ ")" + +type HakaruProxy = ('KProxy :: KProxy Hakaru) +type LiveSet = VarSet HakaruProxy +type HakaruVar = SomeVariable HakaruProxy + +-- The @HoistM@ monad makes use of three monadic layers to propagate information +-- both downwards to the leaves and upwards to the root node of the AST. +-- +-- The Writer layer propagates the live expressions which may be hoisted (i.e. +-- all their data dependencies are currently filled) from each subexpression to +-- their parents. +-- +-- The Reader layer propagates the currently bound variables which will be used +-- to decide when to introduce new bindings. +-- +-- The State layer is just to provide a counter in order to gensym new +-- variables, since the process of adding new bindings is a little tricky. +-- What we want is to fully duplicate bindings without altering the original +-- variable identifiers. To do so, all original variable names are preserved and +-- new variables are added outside the range of existing variables. +newtype HoistM (abt :: [Hakaru] -> Hakaru -> *) a + = HoistM { runHoistM :: RWS LiveSet (ExpressionSet abt) Nat a } + +deriving instance Functor (HoistM abt) +deriving instance (ABT Term abt) => Applicative (HoistM abt) +deriving instance (ABT Term abt) => Monad (HoistM abt) +deriving instance (ABT Term abt) => MonadState Nat (HoistM abt) +deriving instance (ABT Term abt) => MonadWriter (ExpressionSet abt) (HoistM abt) +deriving instance (ABT Term abt) => MonadReader LiveSet (HoistM abt) + +newtype ExpressionSet (abt :: [Hakaru] -> Hakaru -> *) + = ExpressionSet [Entry (abt '[])] + +mergeEntry :: (ABT Term abt) => Entry (abt '[]) -> Entry (abt '[]) -> Entry (abt '[]) +mergeEntry (Entry d e s1 b1) (Entry _ _ s2 b2) = + case jmEq1 s1 s2 of + Just Refl -> Entry d e s1 $ L.nub (b1 ++ b2) + Nothing -> error "cannot union mismatched entries" + +entryEqual :: (ABT Term abt) => Entry (abt '[]) -> Entry (abt '[]) -> Bool +entryEqual Entry{varDependencies=d1,expression=e1,sing=s1} + Entry{varDependencies=d2,expression=e2,sing=s2} = + case (d1 == d2, jmEq1 s1 s2) of + (True , Just Refl) -> alphaEq e1 e2 + _ -> False + +unionEntrySet + :: forall abt + . (ABT Term abt) + => ExpressionSet abt + -> ExpressionSet abt + -> ExpressionSet abt +unionEntrySet (ExpressionSet xs) (ExpressionSet ys) = + ExpressionSet . mapMaybe uniquify $ L.groupBy entryEqual (xs ++ ys) + where + uniquify :: [Entry (abt '[])] -> Maybe (Entry (abt '[])) + uniquify [] = Nothing + uniquify zs = Just $ L.foldl1' mergeEntry zs + +intersectEntrySet + :: forall abt + . (ABT Term abt) + => ExpressionSet abt + -> ExpressionSet abt + -> ExpressionSet abt +intersectEntrySet (ExpressionSet xs) (ExpressionSet ys) = ExpressionSet merged + where + merged :: [Entry (abt '[])] + merged = map (uncurry mergeEntry) . filter (uncurry entryEqual) $ liftA2 (,) xs ys + +-- The general case for generating the entry set for a term is to simply union +-- the sets for all the subterms, so we choose union as our monoidal operation +-- for the Writer monad. +instance (ABT Term abt) => Monoid (ExpressionSet abt) where + mempty = ExpressionSet [] + mappend = unionEntrySet + +-- Given a list of entries to introduce, order them so that their data +-- data dependencies are satisified. +topSortEntries + :: forall abt + . [Entry (abt '[])] + -> [Entry (abt '[])] +topSortEntries entryList = map (entries V.!) $ G.topSort graph + where + entries :: V.Vector (Entry (abt '[])) + !entries = V.fromList entryList + + -- The graph is represented as dependencies between entries, where an entry + -- (a) depends on entry (b) if (b) introduces a variable which (a) depends + -- on. + getVIDs :: Entry (abt '[]) -> [Int] + getVIDs Entry{bindings=b} = map (fromNat . varID) b + + -- Associates all variables introduced by an entry to the entry itself. + -- A given entry may introduce multiple bindings, since an entry stores all + -- α-equivalent variable definitions. + assocBindingsTo :: IM.IntMap Int -> Int -> Entry (abt '[]) -> IM.IntMap Int + assocBindingsTo m n = L.foldl' (\acc v -> IM.insert v n acc) m . getVIDs + + -- Mapping from variable IDs to their corresponding entries + varMap :: IM.IntMap Int + !varMap = V.ifoldl' assocBindingsTo IM.empty entries + + -- Create an edge from each dependency to the variable + makeEdges :: Int -> Entry (abt '[]) -> [G.Edge] + makeEdges idx Entry{varDependencies=d} = map (, idx) + . mapMaybe (flip IM.lookup varMap) + $ varSetKeys d + + -- Collect all the verticies to build the full graph + vertices :: [G.Edge] + !vertices = V.foldr (++) [] $ V.imap makeEdges entries + + -- The full graph structure to be topologically sorted + graph :: G.Graph + !graph = G.buildG (0, V.length entries - 1) vertices + +recordEntry + :: (ABT Term abt) + => Variable a + -> abt '[] a + -> HoistM abt () +recordEntry v abt = tell $ ExpressionSet [Entry (freeVars abt) abt (varType v) [v]] + +execHoistM :: Nat -> HoistM abt a -> a +execHoistM counter act = a + where + hoisted = runHoistM act + (a, _, _) = runRWS hoisted emptyVarSet counter + +-- | An expression is considered "toplevel" if it can be hoisted outside all +-- binders. This means that the expression has no data dependencies. +toplevelEntry + :: Entry abt + -> Bool +toplevelEntry Entry{varDependencies=d} = sizeVarSet d == 0 + +captureEntries + :: (ABT Term abt) + => HoistM abt a + -> HoistM abt (a, ExpressionSet abt) +captureEntries = censor (const mempty) . listen + +hoist + :: (ABT Term abt) + => abt '[] a + -> abt '[] a +hoist abt = execHoistM (nextFreeOrBind abt) $ + captureEntries (hoist' abt) >>= uncurry (introduceToplevel emptyVarSet) + +partitionEntrySet + :: (Entry (abt '[]) -> Bool) + -> ExpressionSet abt + -> (ExpressionSet abt, ExpressionSet abt) +partitionEntrySet p (ExpressionSet xs) = (ExpressionSet true, ExpressionSet false) + where + (true, false) = L.partition p xs + +introduceToplevel + :: (ABT Term abt) + => LiveSet + -> abt '[] a + -> ExpressionSet abt + -> HoistM abt (abt '[] a) +introduceToplevel avail abt entries = do + -- After transforming the given ast, we need to introduce all the toplevel + -- bindings (i.e. bindings with no data dependencies), most of which should be + -- eliminated by constant propagation. + let (ExpressionSet toplevel, rest) = partitionEntrySet toplevelEntry entries + intro = concatMap getBoundVars toplevel ++ fromVarSet avail + -- First we wrap the now AST in the all terms which depdend on top level + -- definitions + wrapped <- introduceBindings intro abt rest + -- Then wrap the result in the toplevel definitions + wrapExpr wrapped toplevel + +bindVar + :: (ABT Term abt) + => Variable (a :: Hakaru) + -> HoistM abt b + -> HoistM abt b +bindVar = local . insertVarSet + +isolateBinder + :: (ABT Term abt) + => Variable (a :: Hakaru) + -> HoistM abt b + -> HoistM abt (b, ExpressionSet abt) +isolateBinder v = captureEntries . bindVar v + +hoist' + :: forall abt xs a . (ABT Term abt) + => abt xs a + -> HoistM abt (abt xs a) +hoist' = start + where + insertMany :: [HakaruVar] -> LiveSet -> LiveSet + insertMany = flip $ L.foldl' (\ acc (SomeVariable v) -> insertVarSet v acc) + + start :: forall ys b . abt ys b -> HoistM abt (abt ys b) + start = loop [] . viewABT + + isolateBinders :: [HakaruVar] -> HoistM abt c -> HoistM abt (c, ExpressionSet abt) + isolateBinders xs = censor (const mempty) . listen . local (insertMany xs) + + -- @loop@ takes 2 parameters. + -- + -- 1. The list of variables bound so far + -- 2. The current term we are recurring over + -- + -- We add a value to the first every time we hit a @Bind@ term, and when + -- a @Syn@ term is finally reached, we introduce any hoisted values whose + -- data dependencies are satisified by these new variables. + loop :: forall ys b + . [HakaruVar] + -> View (Term abt) ys b + -> HoistM abt (abt ys b) + loop _ (Var v) = return (var v) + + -- This case is not needed, but we can avoid performing the expensive work + -- of calling introduceBindings in the case were we won't be performing any + -- work. + loop [] (Syn s) = hoistTerm s + loop xs (Syn s) = do + (term, entries) <- isolateBinders xs (hoistTerm s) + introduceBindings xs term entries + + loop xs (Bind v b) = bind v <$> loop (SomeVariable v : xs) b + +getBoundVars :: Entry x -> [HakaruVar] +getBoundVars Entry{bindings=b} = fmap SomeVariable b + +wrapExpr + :: forall abt b . (ABT Term abt) + => abt '[] b + -> [Entry (abt '[])] + -> HoistM abt (abt '[] b) +wrapExpr = F.foldrM wrap + where + mklet :: abt '[] a -> Variable a -> abt '[] b -> abt '[] b + mklet e v b = + case viewABT b of + Var v' | Just Refl <- varEq v v' -> e + _ -> syn (Let_ :$ e :* bind v b :* End) + + -- Binds the Entry's expression to a fresh variable and rebinds any other + -- variable uses to the fresh variable. + wrap :: Entry (abt '[]) -> abt '[] b -> HoistM abt (abt '[] b) + wrap Entry{expression=e,bindings=[]} acc = do + tmp <- varForExpr e + return $ mklet e tmp acc + wrap Entry{expression=e,bindings=(x:xs)} acc = do + let rhs = var x + body = foldr (mklet rhs) acc xs + return $ mklet e x body + +-- This will introduce all binders which must be introduced by binding the +-- @newVars@ set. As a side effect, the remaining entries are written into the +-- Writer layer of the stack. +introduceBindings + :: forall (a :: Hakaru) abt + . (ABT Term abt) + => [HakaruVar] + -> abt '[] a + -> ExpressionSet abt + -> HoistM abt (abt '[] a) +introduceBindings newVars body (ExpressionSet entries) = do + tell (ExpressionSet leftOver) + wrapExpr body (topSortEntries resultEntries) + where + resultEntries, leftOver :: [Entry (abt '[])] + (resultEntries, leftOver) = loop entries newVars + + introducedBy + :: forall (b :: Hakaru) + . Variable b + -> Entry (abt '[]) + -> Bool + introducedBy v Entry{varDependencies=deps} = memberVarSet v deps + + loop + :: [Entry (abt '[])] + -> [HakaruVar] + -> ([Entry (abt '[])], [Entry (abt '[])]) + loop exprs [] = ([], exprs) + loop exprs (SomeVariable v : xs) = (introduced ++ intro, acc) + where + ~(intro, acc) = loop rest (xs ++ vars) + vars = concatMap getBoundVars introduced + (introduced, rest) = L.partition (introducedBy v) exprs + +-- Contrary to the other binding forms, let expressions are killed by the +-- hoisting pass. Their RHSs are floated upward in the AST and re-introduced +-- where their data dependencies are fulfilled. Thus, the result of hoisting +-- a let expression is just the hoisted body. +hoistTerm + :: forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *) + . (ABT Term abt) + => Term abt a + -> HoistM abt (abt '[] a) +hoistTerm (Let_ :$ rhs :* body :* End) = + caseBind body $ \ v body' -> do + rhs' <- hoist' rhs + recordEntry v rhs' + bindVar v (hoist' body') + +hoistTerm (Lam_ :$ body :* End) = + caseBind body $ \ v body' -> do + available <- fmap (insertVarSet v) ask + (body'', entries) <- isolateBinder v (hoist' body') + finalized <- introduceToplevel available body'' entries + return $ syn (Lam_ :$ bind v finalized :* End) + +hoistTerm term = do + result <- syn <$> traverse21 hoist' term + if isValue result + then return result + else do fresh <- varForExpr result + recordEntry fresh result + return (var fresh) + diff --git a/haskell/Language/Hakaru/Syntax/IClasses.hs b/haskell/Language/Hakaru/Syntax/IClasses.hs --- a/haskell/Language/Hakaru/Syntax/IClasses.hs +++ b/haskell/Language/Hakaru/Syntax/IClasses.hs @@ -25,7 +25,7 @@ -- TODO: DeriveDataTypeable for all our newtypes? ---------------------------------------------------------------- module Language.Hakaru.Syntax.IClasses - ( + ( -- * Showing indexed types Show1(..), shows1, showList1 , Show2(..), shows2, showList2 @@ -58,12 +58,14 @@ , Functor21(..) , Functor22(..) , Foldable11(..), Lift1(..) + , Foldable12(..) , Foldable21(..), Lift2(..) , Foldable22(..) , Traversable11(..) + , Traversable12(..) , Traversable21(..) , Traversable22(..) - + -- * Helper types , Some1(..) , Some2(..) @@ -200,7 +202,7 @@ . showString " " . showsPrec1 11 e2 ) - + showParen_02 :: (Show b, Show2 a) => Int -> String -> b -> a i j -> ShowS showParen_02 p s e1 e2 = showParen (p > 9) @@ -367,7 +369,10 @@ ---------------------------------------------------------------- ---------------------------------------------------------------- --- TODO: rather than having this plethora of classes for different indexing, define newtypes for 1-natural transformations, 2-natural transformations, etc; and then define a single higher-order functor class which is parameterized by the input and output categories. +-- TODO: rather than having this plethora of classes for different +-- indexing, define newtypes for 1-natural transformations, 2-natural +-- transformations, etc; and then define a single higher-order functor +-- class which is parameterized by the input and output categories. -- | A functor on the category of @k@-indexed types (i.e., from -- @k@-indexed types to @k@-indexed types). We unify the two indices, @@ -379,7 +384,7 @@ -- we can derive for 'Functor'. class Functor11 (f :: (k1 -> *) -> k2 -> *) where fmap11 :: (forall i. a i -> b i) -> f a j -> f b j - + class Functor12 (f :: (k1 -> *) -> k2 -> k3 -> *) where fmap12 :: (forall i. a i -> b i) -> f a j l -> f b j l @@ -454,6 +459,16 @@ -- TODO: standard Foldable wrappers 'and11', 'or11', 'all11', 'any11',... +class Functor12 f => Foldable12 (f :: (k1 -> *) -> k2 -> k3 -> *) where + {-# MINIMAL fold12 | foldMap12 #-} + + fold12 :: (Monoid m) => f (Lift1 m) j l -> m + fold12 = foldMap12 unLift1 + + foldMap12 :: (Monoid m) => (forall i. a i -> m) -> f a j l -> m + foldMap12 f = fold12 . fmap12 (Lift1 . f) + + class Functor21 f => Foldable21 (f :: (k1 -> k2 -> *) -> k3 -> *) where {-# MINIMAL fold21 | foldMap21 #-} @@ -490,6 +505,12 @@ sequence :: Monad m => t (m a) -> m (t a) -} +class Foldable12 t => Traversable12 (t :: (k1 -> *) -> k2 -> k3 -> *) where + traverse12 + :: Applicative f + => (forall i. a i -> f (b i)) + -> t a j l + -> f (t b j l) class Foldable21 t => Traversable21 (t :: (k1 -> k2 -> *) -> k3 -> *) where traverse21 @@ -497,7 +518,6 @@ => (forall h i. a h i -> f (b h i)) -> t a j -> f (t b j) - class Foldable22 t => Traversable22 (t :: (k1 -> k2 -> *) -> k3 -> k4 -> *) @@ -508,7 +528,6 @@ -> t a j l -> f (t b j l) - ---------------------------------------------------------------- ---------------------------------------------------------------- -- | Any unindexed type can be lifted to be (trivially) @k@-indexed. @@ -537,10 +556,10 @@ instance Show a => Show1 (Lift2 a i) where showsPrec1 p (Lift2 x) = showsPrec p x show1 (Lift2 x) = show x - + instance Eq a => Eq2 (Lift2 a) where eq2 (Lift2 a) (Lift2 b) = a == b - + instance Eq a => Eq1 (Lift2 a i) where eq1 (Lift2 a) (Lift2 b) = a == b @@ -626,7 +645,7 @@ -- BUG: how do we actually use the term-level @(++)@ at the type level? Or do we have to redefine it ourselves (as below)? If we define it ourselves, how can we make sure it's usable? In particular, how can we prove associativity and that @'[]@ is a /two-sided/ identity element? type family (xs :: [k]) ++ (ys :: [k]) :: [k] where '[] ++ ys = ys - (x ': xs) ++ ys = x ': (xs ++ ys) + (x ': xs) ++ ys = x ': (xs ++ ys) {- -- BUG: having the instances for @[[HakaruFun]]@ and @[HakaruFun]@ precludes giving a general kind-polymorphic data instance for type-level lists; so we have to monomorphize it to just the @[Hakaru]@ kind. diff --git a/haskell/Language/Hakaru/Syntax/Prelude.hs b/haskell/Language/Hakaru/Syntax/Prelude.hs --- a/haskell/Language/Hakaru/Syntax/Prelude.hs +++ b/haskell/Language/Hakaru/Syntax/Prelude.hs @@ -113,11 +113,11 @@ , nil, cons, list -- * Lambda calculus - , lam, lamWithVar, let_ + , lam, lamWithVar, let_, letM , app, app2, app3 -- * Arrays - , empty, arrayWithVar, array, (!), size, reduce + , empty, arrayWithVar, array, arrayLit, (!), size, reduce , sumV, summateV, appendV, mapV, mapWithIndex, normalizeV, constV, unitV, zipWithV -- * Implementation details @@ -125,10 +125,14 @@ , arrayOp0_, arrayOp1_, arrayOp2_, arrayOp3_ , measure0_, measure1_, measure2_ , unsafeNaryOp_, naryOp_withIdentity, naryOp2_ + + -- * Reducers + , bucket, r_fanout, r_index, r_split, r_nop, r_add + ) where -- TODO: implement and use Prelude's fromInteger and fromRational, so we can use numeric literals! -import Prelude (Maybe(..), Bool(..), Integer, Rational, ($), flip, const, error) +import Prelude (Maybe(..), Functor(..), Bool(..), Integer, Rational, ($), flip, const, error) import qualified Prelude import Data.Sequence (Seq) import qualified Data.Sequence as Seq @@ -137,6 +141,8 @@ import qualified Data.List.NonEmpty as L import Data.Semigroup (Semigroup(..)) import Control.Category (Category(..)) +import Control.Monad (return) +import Control.Monad.Fix import Data.Number.Natural import Language.Hakaru.Types.DataKind @@ -144,6 +150,7 @@ import Language.Hakaru.Syntax.TypeOf import Language.Hakaru.Types.HClasses import Language.Hakaru.Types.Coercion +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.Datum import Language.Hakaru.Syntax.ABT hiding (View(..)) @@ -961,6 +968,12 @@ -> abt '[] b let_ e f = syn (Let_ :$ e :* binder Text.empty (typeOf e) f :* End) +letM :: (Functor m, MonadFix m, ABT Term abt) + => abt '[] a + -> (abt '[] a -> m (abt '[] b)) + -> m (abt '[] b) +letM e f = fmap (\ body -> syn $ Let_ :$ e :* body :* End) (binderM Text.empty t f) + where t = typeOf e ---------------------------------------------------------------- array @@ -969,7 +982,7 @@ -> (abt '[] 'HNat -> abt '[] a) -> abt '[] ('HArray a) array n = - syn . Array_ n . binder Text.empty sing + syn . Array_ n . binder Text.empty sing arrayWithVar :: (ABT Term abt) @@ -980,6 +993,11 @@ arrayWithVar n x body = syn $ Array_ n (bind x body) +arrayLit + :: (ABT Term abt) + => [abt '[] a] + -> abt '[] ('HArray a) +arrayLit = syn . ArrayLiteral_ empty :: (ABT Term abt, SingI a) => abt '[] ('HArray a) empty = syn (Empty_ sing) @@ -1135,7 +1153,48 @@ zipWithV f v1 v2 = array (size v1) (\i -> f (v1 ! i) (v2 ! i)) +---------------------------------------------------------------- +r_fanout + :: (ABT Term abt) + => Reducer abt xs a + -> Reducer abt xs b + -> Reducer abt xs (HPair a b) +r_fanout = Red_Fanout + +r_index + :: (Binders Term abt xs as) + => (as -> abt '[] 'HNat) + -> ((abt '[] 'HNat, as) -> abt '[] 'HNat) + -> Reducer abt ( 'HNat ': xs) a + -> Reducer abt xs ('HArray a) +r_index n f = Red_Index (binders n) (binders f) + +r_split + :: (Binders Term abt xs as) + => ((abt '[] 'HNat, as) -> abt '[] HBool) + -> Reducer abt xs a + -> Reducer abt xs b + -> Reducer abt xs (HPair a b) +r_split b = Red_Split (binders b) + +r_nop :: (ABT Term abt) => Reducer abt xs HUnit +r_nop = Red_Nop + +r_add + :: (Binders Term abt xs as, HSemiring_ a) + => ((abt '[] 'HNat, as) -> abt '[] a) + -> Reducer abt xs a +r_add f = Red_Add hSemiring (binders f) + +bucket + :: (ABT Term abt) + => abt '[] 'HNat + -> abt '[] 'HNat + -> Reducer abt '[] a + -> abt '[] a +bucket i j r = syn $ Bucket i j r + ---------------------------------------------------------------- (>>=) :: (ABT Term abt) @@ -1584,10 +1643,9 @@ exponential (prob_ 1) >>= \x -> dirac $ b * x ** recip k --- BUG: would it be better to 'observe' that @p <= 1@ before doing the superpose? At least that way things would be /defined/ for all inputs... bern :: (ABT Term abt) => abt '[] 'HProb -> abt '[] ('HMeasure HBool) -bern p = weightedDirac true p - <|> weightedDirac false (prob_ 1 `unsafeMinusProb` p) +bern p = categorical (arrayLit [p, prob_ 1 `unsafeMinusProb` p]) >>= \i -> + dirac (arrayLit [true, false] ! i) mix :: (ABT Term abt) => abt '[] ('HArray 'HProb) -> abt '[] ('HMeasure 'HNat) diff --git a/haskell/Language/Hakaru/Syntax/Prune.hs b/haskell/Language/Hakaru/Syntax/Prune.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Prune.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE CPP + , DataKinds + , FlexibleContexts + , GADTs + , GeneralizedNewtypeDeriving + , KindSignatures + , MultiParamTypeClasses + , ScopedTypeVariables + , TypeOperators + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.Prune +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Prune (prune) where + +import Control.Monad.Reader +import Data.Maybe + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Syntax.Unroll (renameInEnv) +import Language.Hakaru.Types.DataKind + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +-- A Simple pass for pruning the unused let bindings from an AST. + +newtype PruneM a = PruneM { runPruneM :: Reader Varmap a } + deriving (Functor, Applicative, Monad, MonadReader Varmap, MonadFix) + +lookupEnv + :: forall (a :: Hakaru) + . Variable a + -> Varmap + -> Variable a +lookupEnv v = fromMaybe v . lookupAssoc v + +prune + :: (ABT Term abt) + => abt '[] a + -> abt '[] a +prune = flip runReader emptyAssocs . runPruneM . prune' + +prune' + :: forall abt xs a . (ABT Term abt) + => abt xs a + -> PruneM (abt xs a) +prune' = loop . viewABT + where + loop :: forall (b :: Hakaru) ys . View (Term abt) ys b -> PruneM (abt ys b) + loop (Var v) = (var . lookupEnv v) `fmap` ask + loop (Syn s) = pruneTerm s + loop (Bind v b) = renameInEnv v (loop b) + +pruneTerm + :: forall a abt + . (ABT Term abt) + => Term abt a + -> PruneM (abt '[] a) +pruneTerm (Let_ :$ rhs :* body :* End) = + caseBind body $ \v body' -> + let frees = freeVars body' + mklet r b = syn (Let_ :$ r :* b :* End) + doRhs = prune' rhs + doBody = prune' body' + fullExpr = mklet <$> doRhs <*> renameInEnv v doBody + in case viewABT body' of + Var v' | Just Refl <- varEq v v' -> doRhs + _ | memberVarSet v frees -> fullExpr + | otherwise -> doBody + +pruneTerm term = syn <$> traverse21 prune' term diff --git a/haskell/Language/Hakaru/Syntax/Reducer.hs b/haskell/Language/Hakaru/Syntax/Reducer.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Reducer.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE CPP + , DataKinds + , InstanceSigs + , GADTs + , KindSignatures + , Rank2Types + , TypeOperators + #-} + +module Language.Hakaru.Syntax.Reducer where + +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.HClasses +import Language.Hakaru.Syntax.IClasses + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +import Data.Monoid (Monoid(..)) +#endif + +data Reducer (abt :: [Hakaru] -> Hakaru -> *) + (xs :: [Hakaru]) + (a :: Hakaru) where + Red_Fanout + :: Reducer abt xs a + -> Reducer abt xs b + -> Reducer abt xs (HPair a b) + Red_Index + :: abt xs 'HNat -- size of resulting array + -> abt ( 'HNat ': xs) 'HNat -- index into array (bound i) + -> Reducer abt ( 'HNat ': xs) a -- reduction body (bound b) + -> Reducer abt xs ('HArray a) + Red_Split + :: abt ( 'HNat ': xs) HBool -- (bound i) + -> Reducer abt xs a + -> Reducer abt xs b + -> Reducer abt xs (HPair a b) + Red_Nop + :: Reducer abt xs HUnit + Red_Add + :: HSemiring a + -> abt ( 'HNat ': xs) a -- (bound i) + -> Reducer abt xs a + +instance Functor22 Reducer where + fmap22 f (Red_Fanout r1 r2) = Red_Fanout (fmap22 f r1) (fmap22 f r2) + fmap22 f (Red_Index n ix r) = Red_Index (f n) (f ix) (fmap22 f r) + fmap22 f (Red_Split b r1 r2) = Red_Split (f b) (fmap22 f r1) (fmap22 f r2) + fmap22 _ Red_Nop = Red_Nop + fmap22 f (Red_Add h e) = Red_Add h (f e) + +instance Foldable22 Reducer where + foldMap22 f (Red_Fanout r1 r2) = foldMap22 f r1 `mappend` foldMap22 f r2 + foldMap22 f (Red_Index n ix r) = f n `mappend` f ix `mappend` foldMap22 f r + foldMap22 f (Red_Split b r1 r2) = f b `mappend` foldMap22 f r1 `mappend` foldMap22 f r2 + foldMap22 _ Red_Nop = mempty + foldMap22 f (Red_Add _ e) = f e + +instance Traversable22 Reducer where + traverse22 f (Red_Fanout r1 r2) = Red_Fanout <$> traverse22 f r1 <*> traverse22 f r2 + traverse22 f (Red_Index n ix r) = Red_Index <$> f n <*> f ix <*> traverse22 f r + traverse22 f (Red_Split b r1 r2) = Red_Split <$> f b <*> traverse22 f r1 <*> traverse22 f r2 + traverse22 f Red_Nop = pure Red_Nop + traverse22 f (Red_Add h e) = Red_Add h <$> f e + + +instance Eq2 abt => Eq1 (Reducer abt xs) where + eq1 (Red_Fanout r1 r2) (Red_Fanout r1' r2') = eq1 r1 r1' && eq1 r2 r2' + eq1 (Red_Index n ix r) (Red_Index n' ix' r') = eq2 n n' && eq2 ix ix' && eq1 r r' + eq1 (Red_Split b r1 r2) (Red_Split b' r1' r2') = eq2 b b' && eq1 r1 r1' && eq1 r2 r2' + eq1 Red_Nop Red_Nop = True + eq1 (Red_Add _ e) (Red_Add _ e') = eq2 e e' + eq1 _ _ = False + +instance JmEq2 abt => JmEq1 (Reducer abt xs) where + jmEq1 = jmEqReducer + +jmEqReducer + :: (JmEq2 abt) + => Reducer abt xs a + -> Reducer abt xs b + -> Maybe (TypeEq a b) +jmEqReducer (Red_Fanout a b) (Red_Fanout a' b') = do + Refl <- jmEqReducer a a' + Refl <- jmEqReducer b b' + return Refl +jmEqReducer (Red_Index s i r) (Red_Index s' i' r') = do + (Refl, Refl) <- jmEq2 s s' + (Refl, Refl) <- jmEq2 i i' + Refl <- jmEqReducer r r' + return Refl +jmEqReducer (Red_Split b r s) (Red_Split b' r' s') = do + (Refl, Refl) <- jmEq2 b b' + Refl <- jmEqReducer r r' + Refl <- jmEqReducer s s' + return Refl +jmEqReducer Red_Nop Red_Nop = return Refl +jmEqReducer (Red_Add _ x) (Red_Add _ x') = do + (Refl, Refl) <- jmEq2 x x' + return Refl +jmEqReducer _ _ = Nothing + diff --git a/haskell/Language/Hakaru/Syntax/Rename.hs b/haskell/Language/Hakaru/Syntax/Rename.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Rename.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE CPP + , DataKinds + , EmptyCase + , ExistentialQuantification + , FlexibleContexts + , GADTs + , GeneralizedNewtypeDeriving + , KindSignatures + , MultiParamTypeClasses + , OverloadedStrings + , PolyKinds + , ScopedTypeVariables + , TypeFamilies + , TypeOperators + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- | +-- Module : Language.Hakaru.Syntax.Rename +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- Performs renaming of variables hints only (in Hakaru expressions) +-- which hopefully has no effect on semantics but can produce prettier expressions +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Rename where + +import Control.Monad.Reader +import Control.Monad.State +import Data.Maybe (fromMaybe) +import Data.Number.Nat + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq (Varmap) +import Language.Hakaru.Syntax.Gensym +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Syntax.Variable +import qualified Data.Text as Text +import Data.Text (Text) +import Data.Char + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +type Renamer = Text -> Text + +renameAST + :: forall abt xs a . (ABT Term abt) + => Renamer + -> abt xs a + -> abt xs a +renameAST r = start + where + start :: abt ys b -> abt ys b + start = loop . viewABT + + loop :: View (Term abt) ys b -> abt ys b + loop (Var v) = var (renameVar r v) + loop (Syn s) = syn (fmap21 start s) + loop (Bind v b) = bind (renameVar r v) (loop b) + +renameVar :: Renamer -> Variable a -> Variable a +renameVar r v = v { varHint = r (varHint v) } + +removeUnicodeChars :: Text -> Text +removeUnicodeChars = Text.filter isAscii diff --git a/haskell/Language/Hakaru/Syntax/TypeCheck.hs b/haskell/Language/Hakaru/Syntax/TypeCheck.hs --- a/haskell/Language/Hakaru/Syntax/TypeCheck.hs +++ b/haskell/Language/Hakaru/Syntax/TypeCheck.hs @@ -68,6 +68,7 @@ , HRadical(..), HContinuous(..)) import Language.Hakaru.Syntax.ABT import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.AST.Sing (sing_Literal, sing_MeasureOp) @@ -164,10 +165,11 @@ -- typing issue. Thus, for non-empty arrays and non-phantom -- record types, we should be able to infer the whole type -- provided we can infer the various subterms. - go U.Empty_ = True - go (U.Pair_ e1 e2) = mustCheck e1 && mustCheck e2 - go (U.Array_ _ e1) = mustCheck' e1 - go (U.Datum_ _) = True + go U.Empty_ = True + go (U.Pair_ e1 e2) = mustCheck e1 && mustCheck e2 + go (U.Array_ _ e1) = mustCheck' e1 + go (U.ArrayLiteral_ es) = F.all mustCheck es + go (U.Datum_ _) = True -- TODO: everyone says this, but it seems to me that if we can -- infer any of the branches (and check the rest to agree) then @@ -184,6 +186,7 @@ go (U.Integrate_ _ _ _) = False go (U.Summate_ _ _ _) = False go (U.Product_ _ _ _) = False + go (U.Bucket_ _ _ _) = False go U.Reject_ = True go (U.Expect_ _ e2) = mustCheck' e2 go (U.Observe_ e1 _) = mustCheck e1 @@ -283,6 +286,7 @@ case (sourceSpan, input_) of (Just s, Just input) -> return $ mconcat [ header + , "\n\n" , U.printSourceSpan s input , footer ] @@ -297,7 +301,7 @@ -> TypeCheckMonad r typeMismatch s typ1 typ2 = failwith =<< makeErrMsg - "Type Mismatch:\n\n" + "Type Mismatch:" s (mconcat [ "expected " , msg1 @@ -315,7 +319,7 @@ -> TypeCheckMonad r missingInstance clas typ s = failwith =<< makeErrMsg - "Missing Instance: " + "Missing Instance:" s (mconcat $ ["No ", clas, " instance for type ", showT typ]) @@ -326,7 +330,7 @@ -> TypeCheckMonad r missingLub typ1 typ2 s = failwith =<< makeErrMsg - "Missing common type:\n\n" + "Missing common type:" s (mconcat ["No lub of types ", showT typ1, " and ", showT typ2]) @@ -458,6 +462,12 @@ Cons1 x xs' -> pushCtx x (bind x <$> checkBinders xs' eTyp e) +-- HACK: Passing this list of variables feels like a hack +-- it would be nice if it could be removed from this datatype +data TypedReducer (abt :: [Hakaru] -> Hakaru -> *) + (xs :: [Hakaru]) + = forall b. TypedReducer !(Sing b) (List1 Variable xs) (Reducer abt xs b) + ---------------------------------------------------------------- -- | Given a typing environment and a term, synthesize the term's -- type (and produce an elaborated term): @@ -598,6 +608,15 @@ inferBinder SNat e2 $ \typ2 e2' -> return . TypedAST (SArray typ2) $ syn (Array_ e1' e2') + U.ArrayLiteral_ es -> do + mode <- getMode + TypedASTs typ es' <- + case mode of + StrictMode -> inferOneCheckOthers_ es + LaxMode -> inferLubType sourceSpan es + UnsafeMode -> inferLubType sourceSpan es + return . TypedAST (SArray typ) $ syn (ArrayLiteral_ es') + U.Case_ e1 branches -> do TypedAST typ1 e1' <- inferType_ e1 mode <- getMode @@ -688,6 +707,13 @@ syn (Product h1 h2 :$ e1' :* e2' :* e3' :* End) _ -> failwith_ "Product given bounds which are not discrete" + U.Bucket_ e1 e2 r1 -> do + e1' <- checkType_ SNat e1 + e2' <- checkType_ SNat e2 + TypedReducer typ1 Nil1 r1' <- inferReducer r1 Nil1 + return . TypedAST typ1 $ + syn (Bucket e1' e2' r1') + U.Expect_ e1 e2 -> do TypedAST typ1 e1' <- inferType_ e1 case typ1 of @@ -948,6 +974,48 @@ _ -> typeMismatch Nothing (Right typ) (Left "HFun") _ -> argumentNumberError + inferReducer :: U.Reducer xs U.U_ABT 'U.U + -> List1 Variable xs1 + -> TypeCheckMonad (TypedReducer abt xs1) + + inferReducer (U.R_Fanout_ r1 r2) xs = do + TypedReducer t1 _ r1' <- inferReducer r1 xs + TypedReducer t2 _ r2' <- inferReducer r2 xs + return (TypedReducer (sPair t1 t2) xs (Red_Fanout r1' r2')) + + inferReducer (U.R_Index_ x n ix r1) xs = do + let (_, n') = caseBinds n + let b = makeVar x SNat + TypedReducer t1 _ r1' <- inferReducer r1 (Cons1 b xs) + n'' <- checkBinders xs SNat n' + caseBind ix $ \i ix1 -> + let i' = makeVar i SNat + (_, ix2) = caseBinds ix1 in do + ix3 <- pushCtx i' (checkBinders xs SNat ix2) + return . TypedReducer (SArray t1) xs $ + Red_Index n'' (bind i' ix3) r1' + + inferReducer (U.R_Split_ b r1 r2) xs = do + TypedReducer t1 _ r1' <- inferReducer r1 xs + TypedReducer t2 _ r2' <- inferReducer r2 xs + caseBind b $ \x b1 -> + let (_, b2) = caseBinds b1 + x' = makeVar x SNat in do + b3 <- pushCtx x' (checkBinders xs sBool b2) + return . TypedReducer (sPair t1 t2) xs $ + (Red_Split (bind x' b3) r1' r2') + + inferReducer U.R_Nop_ xs = return (TypedReducer sUnit xs Red_Nop) + + inferReducer (U.R_Add_ e) xs = + caseBind e $ \x e1 -> + let (_, e2) = caseBinds e1 + x' = makeVar x SNat in + pushCtx x' $ + inferBinders xs e2 $ \typ e3 -> do + h <- getHSemiring typ + return $ TypedReducer typ xs (Red_Add h (bind x' e3)) + make_NaryOp :: Sing a -> U.NaryOp -> TypeCheckMonad (NaryOp a) make_NaryOp a U.And = isBool a >>= \Refl -> return And make_NaryOp a U.Or = isBool a >>= \Refl -> return Or @@ -1029,7 +1097,6 @@ showParen_1x p "TypedASTs" typ es -} - -- TODO: can we make this lazier in the second component of 'TypedASTs' -- so that we can perform case analysis on the type component before -- actually evaluating 'checkOthers'? Problem is, even though we @@ -1361,6 +1428,13 @@ e1' <- checkType_ SNat e1 e2' <- checkBinder SNat typ1 e2 return $ syn (Array_ e1' e2') + _ -> typeMismatch sourceSpan (Right typ0) (Left "HArray") + + U.ArrayLiteral_ es -> + case typ0 of + SArray typ1 -> do + es' <- T.forM es $ checkType_ typ1 + return $ syn (ArrayLiteral_ es') _ -> typeMismatch sourceSpan (Right typ0) (Left "HArray") U.Datum_ (U.Datum hint d) -> diff --git a/haskell/Language/Hakaru/Syntax/TypeOf.hs b/haskell/Language/Hakaru/Syntax/TypeOf.hs --- a/haskell/Language/Hakaru/Syntax/TypeOf.hs +++ b/haskell/Language/Hakaru/Syntax/TypeOf.hs @@ -25,6 +25,7 @@ ( -- * Get singletons for well-typed ABTs typeOf + , typeOfReducer -- * Implementation details , getTermSing @@ -32,7 +33,7 @@ import qualified Data.Foldable as F #if __GLASGOW_HASKELL__ < 710 -import Data.Functor ((<$>)) +import Control.Applicative (Applicative(..), (<$>)) #endif import Language.Hakaru.Syntax.IClasses (Pair2(..), fst2, snd2) @@ -40,10 +41,11 @@ import Language.Hakaru.Syntax.ABT (ABT, caseBind, paraABT) import Language.Hakaru.Types.DataKind (Hakaru()) import Language.Hakaru.Types.HClasses (sing_HSemiring) -import Language.Hakaru.Types.Sing (Sing(..), sUnMeasure) +import Language.Hakaru.Types.Sing (Sing(..), sUnMeasure, sUnit, sPair) import Language.Hakaru.Types.Coercion (singCoerceCod, singCoerceDom, Coerce(..)) import Language.Hakaru.Syntax.Datum (Datum(..), Branch(..)) +import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.AST (Term(..), SCon(..), SArgs(..)) import Language.Hakaru.Syntax.AST.Sing (sing_PrimOp, sing_ArrayOp, sing_MeasureOp, sing_NaryOp, sing_Literal) @@ -92,6 +94,16 @@ (LiftSing . getTermSing unLiftSing) + +typeOfReducer + :: Reducer abt xs a + -> Sing a +typeOfReducer (Red_Fanout a b) = sPair (typeOfReducer a) (typeOfReducer b) +typeOfReducer (Red_Index _ _ a) = SArray (typeOfReducer a) +typeOfReducer (Red_Split _ a b) = sPair (typeOfReducer a) (typeOfReducer b) +typeOfReducer Red_Nop = sUnit +typeOfReducer (Red_Add h _) = sing_HSemiring h + -- | This newtype serves two roles. First we add the phantom @xs@ -- so that we can fit this in with the types of 'paraABT'. And -- second, we wrap up the 'Sing' in a monad for capturing errors, @@ -169,6 +181,8 @@ go (Literal_ v) = return $ sing_Literal v go (Empty_ typ) = return typ go (Array_ _ r2) = SArray <$> getSing r2 + go (ArrayLiteral_ es) = SArray <$> tryAll "ArrayLiteral_" getSing es + go (Bucket _ _ r) = return (typeOfReducer r) go (Datum_ (Datum _ typ _)) = return typ go (Case_ _ bs) = tryAll "Case_" getBranchSing bs go (Superpose_ pes) = tryAll "Superpose_" (getSing . snd) pes diff --git a/haskell/Language/Hakaru/Syntax/Uniquify.hs b/haskell/Language/Hakaru/Syntax/Uniquify.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Uniquify.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE CPP + , DataKinds + , EmptyCase + , ExistentialQuantification + , FlexibleContexts + , GADTs + , GeneralizedNewtypeDeriving + , KindSignatures + , MultiParamTypeClasses + , OverloadedStrings + , PolyKinds + , ScopedTypeVariables + , TypeFamilies + , TypeOperators + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.Uniquify +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- Performs renaming of Hakaru expressions to ensure globally unique variable +-- identifiers. +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Uniquify (uniquify) where + +import Control.Monad.Reader +import Control.Monad.State +import Data.Maybe (fromMaybe) +import Data.Number.Nat + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq (Varmap) +import Language.Hakaru.Syntax.Gensym +import Language.Hakaru.Syntax.IClasses +import Debug.Trace + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +newtype Uniquifier a = Uniquifier { runUniquifier :: StateT Nat (Reader Varmap) a } + deriving (Functor, Applicative, Monad, MonadState Nat, MonadReader Varmap) + +uniquify :: (ABT Term abt) => abt '[] a -> abt '[] a +uniquify abt = fst $ runReader (runStateT unique seed) emptyAssocs + where + unique = runUniquifier (uniquify' abt) + seed = nextFreeOrBind abt + +uniquify' + :: forall abt xs a . (ABT Term abt) + => abt xs a + -> Uniquifier (abt xs a) +uniquify' = start + where + start :: abt ys b -> Uniquifier (abt ys b) + start = loop . viewABT + + loop :: View (Term abt) ys b -> Uniquifier (abt ys b) + loop (Var v) = uniquifyVar v + loop (Syn s) = fmap syn (traverse21 start s) + loop (Bind v b) = do + fresh <- freshVar v + let assoc = Assoc v fresh + -- Process the body with the updated Varmap and wrap the + -- result in a bind form + bind fresh <$> local (insertAssoc assoc) (loop b) + +uniquifyVar + :: (ABT Term abt) + => Variable a + -> Uniquifier (abt '[] a) +uniquifyVar v = (var . fromMaybe v . lookupAssoc v) <$> ask + diff --git a/haskell/Language/Hakaru/Syntax/Unroll.hs b/haskell/Language/Hakaru/Syntax/Unroll.hs new file mode 100644 --- /dev/null +++ b/haskell/Language/Hakaru/Syntax/Unroll.hs @@ -0,0 +1,159 @@ +{-# LANGUAGE CPP + , DataKinds + , FlexibleContexts + , GADTs + , GeneralizedNewtypeDeriving + , MultiParamTypeClasses + , RankNTypes + , ScopedTypeVariables + , TypeOperators + #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2017.02.01 +-- | +-- Module : Language.Hakaru.Syntax.Unroll +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : +-- Stability : experimental +-- Portability : GHC-only +-- +-- Performs renaming of Hakaru expressions to ensure globally unique variable +-- identifiers. +-- +---------------------------------------------------------------- +module Language.Hakaru.Syntax.Unroll (renameInEnv, unroll) where + +import Control.Monad.Reader +import Data.Maybe (fromMaybe) +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq (Varmap) +import Language.Hakaru.Syntax.Prelude hiding ((>>=)) +import Language.Hakaru.Types.HClasses +import Prelude hiding (product, (*), (+), (-), + (==), (>=)) + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif + +newtype Unroll a = Unroll { runUnroll :: Reader Varmap a } + deriving (Functor, Applicative, Monad, MonadReader Varmap, MonadFix) + +rebind + :: (ABT Term abt, MonadFix m) + => Variable a + -> (Variable a -> m (abt xs b)) + -> m (abt (a ': xs) b) +rebind source f = binderM (varHint source) (varType source) $ \var' -> + let v = caseVarSyn var' id (const $ error "oops") + in f v + +renameInEnv + :: (ABT Term abt, MonadReader Varmap m, MonadFix m) + => Variable a + -> m (abt xs b) + -> m (abt (a ': xs) b) +renameInEnv source action = rebind source $ \v -> + local (insertAssoc $ Assoc source v) action + +unroll :: forall abt xs a . (ABT Term abt) => abt xs a -> abt xs a +unroll abt = runReader (runUnroll $ unroll' abt) emptyAssocs + +unroll' :: forall abt xs a . (ABT Term abt) => abt xs a -> Unroll (abt xs a) +unroll' = cataABTM var_ renameInEnv (>>= unrollTerm) + where + var_ :: Variable b -> Unroll (abt '[] b) + var_ v = fmap (var . fromMaybe v . lookupAssoc v) ask + +mklet :: ABT Term abt => abt '[] b -> abt '[b] a -> abt '[] a +mklet rhs body = syn (Let_ :$ rhs :* body :* End) + +mksummate, mkproduct + :: (ABT Term abt) + => HDiscrete a + -> HSemiring b + -> abt '[] a + -> abt '[] a + -> abt '[a] b + -> abt '[] b +mksummate a b lo hi body = syn (Summate a b :$ lo :* hi :* body :* End) +mkproduct a b lo hi body = syn (Product a b :$ lo :* hi :* body :* End) + +unrollTerm + :: (ABT Term abt) + => Term abt a + -> Unroll (abt '[] a) +unrollTerm (Summate disc semi :$ lo :* hi :* body :* End) = + case (disc, semi) of + (HDiscrete_Nat, HSemiring_Nat) -> unrollSummate disc semi lo hi body + (HDiscrete_Nat, HSemiring_Int) -> unrollSummate disc semi lo hi body + (HDiscrete_Nat, HSemiring_Prob) -> unrollSummate disc semi lo hi body + (HDiscrete_Nat, HSemiring_Real) -> unrollSummate disc semi lo hi body + + (HDiscrete_Int, HSemiring_Nat) -> unrollSummate disc semi lo hi body + (HDiscrete_Int, HSemiring_Int) -> unrollSummate disc semi lo hi body + (HDiscrete_Int, HSemiring_Prob) -> unrollSummate disc semi lo hi body + (HDiscrete_Int, HSemiring_Real) -> unrollSummate disc semi lo hi body + +unrollTerm (Product disc semi :$ lo :* hi :* body :* End) = + case (disc, semi) of + (HDiscrete_Nat, HSemiring_Nat) -> unrollProduct disc semi lo hi body + (HDiscrete_Nat, HSemiring_Int) -> unrollProduct disc semi lo hi body + (HDiscrete_Nat, HSemiring_Prob) -> unrollProduct disc semi lo hi body + (HDiscrete_Nat, HSemiring_Real) -> unrollProduct disc semi lo hi body + + (HDiscrete_Int, HSemiring_Nat) -> unrollProduct disc semi lo hi body + (HDiscrete_Int, HSemiring_Int) -> unrollProduct disc semi lo hi body + (HDiscrete_Int, HSemiring_Prob) -> unrollProduct disc semi lo hi body + (HDiscrete_Int, HSemiring_Real) -> unrollProduct disc semi lo hi body + +unrollTerm term = return (syn term) + +-- Conditionally introduce a variable for the rhs if the rhs is not currently a +-- variable already. Be careful that the provided variable has been remaped to +-- its equivalent in the target term if altering the binding structure of the +-- program. +letM' :: (Functor m, MonadFix m, ABT Term abt) + => abt '[] a + -> (abt '[] a -> m (abt '[] b)) + -> m (abt '[] b) +letM' e f = + case viewABT e of + Var _ -> f e + Syn (Literal_ _) -> f e + _ -> letM e f + +unrollSummate + :: (ABT Term abt, HSemiring_ a, HSemiring_ b, HEq_ a) + => HDiscrete a + -> HSemiring b + -> abt '[] a + -> abt '[] a + -> abt '[a] b + -> Unroll (abt '[] b) +unrollSummate disc semi lo hi body = + letM' lo $ \loVar -> + letM' hi $ \hiVar -> + let preamble = mklet loVar body + loop = mksummate disc semi (loVar + one) hiVar body + in return $ if_ (loVar == hiVar) zero (preamble + loop) + +unrollProduct + :: (ABT Term abt, HSemiring_ a, HSemiring_ b, HEq_ a) + => HDiscrete a + -> HSemiring b + -> abt '[] a + -> abt '[] a + -> abt '[a] b + -> Unroll (abt '[] b) +unrollProduct disc semi lo hi body = + letM' lo $ \loVar -> + letM' hi $ \hiVar -> + let preamble = mklet loVar body + loop = mkproduct disc semi (loVar + one) hiVar body + in return $ if_ (loVar == hiVar) one (preamble * loop) + diff --git a/haskell/Language/Hakaru/Syntax/Value.hs b/haskell/Language/Hakaru/Syntax/Value.hs --- a/haskell/Language/Hakaru/Syntax/Value.hs +++ b/haskell/Language/Hakaru/Syntax/Value.hs @@ -14,7 +14,10 @@ import Language.Hakaru.Types.HClasses import Language.Hakaru.Types.DataKind import Language.Hakaru.Types.Coercion +import Language.Hakaru.Types.Sing +import Data.STRef + import qualified Data.Vector as V import qualified Data.Number.LogFloat as LF import Data.Number.Nat @@ -41,11 +44,13 @@ VArray :: {-# UNPACK #-} !(V.Vector (Value a)) -> Value ('HArray a) instance Eq1 Value where - eq1 (VNat a) (VNat b) = a == b - eq1 (VInt a) (VInt b) = a == b - eq1 (VProb a) (VProb b) = a == b - eq1 (VReal a) (VReal b) = a == b - eq1 _ _ = False + eq1 (VNat a) (VNat b) = a == b + eq1 (VInt a) (VInt b) = a == b + eq1 (VProb a) (VProb b) = a == b + eq1 (VReal a) (VReal b) = a == b + eq1 (VDatum a) (VDatum b) = a == b + eq1 (VArray a) (VArray b) = a == b + eq1 _ _ = False instance Eq (Value a) where (==) = eq1 @@ -104,3 +109,15 @@ -> [Value a] enumFromUntilValue _ (VNat lo) (VNat hi) = map VNat (init (enumFromTo lo hi)) enumFromUntilValue _ (VInt lo) (VInt hi) = map VInt (init (enumFromTo lo hi)) + +data VReducer :: * -> Hakaru -> * where + VRed_Num :: STRef s (Value a) + -> VReducer s a + VRed_Unit :: VReducer s HUnit + VRed_Pair :: Sing a + -> Sing b + -> VReducer s a + -> VReducer s b + -> VReducer s (HPair a b) + VRed_Array :: V.Vector (VReducer s a) + -> VReducer s ('HArray a) diff --git a/haskell/Language/Hakaru/Syntax/Variable.hs b/haskell/Language/Hakaru/Syntax/Variable.hs --- a/haskell/Language/Hakaru/Syntax/Variable.hs +++ b/haskell/Language/Hakaru/Syntax/Variable.hs @@ -6,6 +6,7 @@ , DeriveDataTypeable , ExistentialQuantification , UndecidableInstances + , ScopedTypeVariables #-} {-# OPTIONS_GHC -Wall -fwarn-tabs #-} @@ -39,9 +40,13 @@ , fromVarSet , toVarSet , toVarSet1 + , varSetKeys , insertVarSet , deleteVarSet , memberVarSet + , unionVarSet + , intersectVarSet + , sizeVarSet , nextVarID -- ** Substitutions; aka: maps from variables to their definitions , Assoc(..) @@ -52,6 +57,7 @@ , toAssocs , toAssocs1 , insertAssoc + , insertOrReplaceAssoc , insertAssocs , lookupAssoc , adjustAssoc @@ -310,6 +316,9 @@ . showsPrec 11 xs ) +instance (Eq (SomeVariable (kproxy :: KProxy k))) => Eq (VarSet kproxy) where + VarSet s1 == VarSet s2 = s1 == s2 + -- | Return the successor of the largest 'varID' of all the variables -- in the set. Thus, we return zero for the empty set and non-zero -- for non-empty sets. @@ -366,12 +375,13 @@ someVariables Nil1 = [] someVariables (Cons1 x xs) = SomeVariable x : someVariables xs - instance Monoid (VarSet kproxy) where mempty = emptyVarSet mappend (VarSet xs) (VarSet ys) = VarSet (IM.union xs ys) -- TODO: remove bias; crash if conflicting definitions mconcat = VarSet . IM.unions . map unVarSet +varSetKeys :: VarSet a -> [Int] +varSetKeys (VarSet set) = IM.keys set insertVarSet :: Variable a -> VarSet (KindOf a) -> VarSet (KindOf a) insertVarSet x (VarSet xs) = @@ -406,6 +416,28 @@ Nothing -> False Just _ -> True +-- NB: The union and intersection operations are left biased. +-- What is the best behaviour when we have two variables with +-- different types in the set? +unionVarSet + :: forall (kproxy :: KProxy k) + . (Show1 (Sing :: k -> *), JmEq1 (Sing :: k -> *)) + => VarSet kproxy + -> VarSet kproxy + -> VarSet kproxy +unionVarSet (VarSet s1) (VarSet s2) = VarSet (IM.union s1 s2) + +intersectVarSet + :: forall (kproxy :: KProxy k) + . (Show1 (Sing :: k -> *), JmEq1 (Sing :: k -> *)) + => VarSet kproxy + -> VarSet kproxy + -> VarSet kproxy +intersectVarSet (VarSet s1) (VarSet s2) = VarSet (IM.intersection s1 s2) + +sizeVarSet :: VarSet a -> Int +sizeVarSet (VarSet xs) = IM.size xs + ---------------------------------------------------------------- -- BUG: haddock doesn't like annotations on GADT constructors. So -- here we'll avoid using the GADT syntax, even though it'd make @@ -515,7 +547,11 @@ insertAssoc v@(Assoc x _) (Assocs xs) = case IM.insertLookupWithKey (\_ v' _ -> v') (fromNat $ varID x) v xs of (Nothing, xs') -> Assocs xs' - (Just _, _) -> error "insertAssoc: variable is already assigned!" + (Just _, _ ) -> error "insertAssoc: variable is already assigned!" + +insertOrReplaceAssoc :: Assoc ast -> Assocs ast -> Assocs ast +insertOrReplaceAssoc v@(Assoc x _) (Assocs xs) = + Assocs $ IM.insert (fromNat $ varID x) v xs insertAssocs :: Assocs ast -> Assocs ast -> Assocs ast insertAssocs (Assocs from) to = IM.foldr insertAssoc to from diff --git a/haskell/Language/Hakaru/Types/Sing.hs b/haskell/Language/Hakaru/Types/Sing.hs --- a/haskell/Language/Hakaru/Types/Sing.hs +++ b/haskell/Language/Hakaru/Types/Sing.hs @@ -6,6 +6,7 @@ , FlexibleInstances , Rank2Types , UndecidableInstances + , ScopedTypeVariables #-} {-# OPTIONS_GHC -Wall -fwarn-tabs #-} @@ -41,6 +42,7 @@ , sUnList , sUnMaybe -- ** Singletons for `Symbol` + , someSSymbol, ssymbolVal , sSymbol_Bool , sSymbol_Unit , sSymbol_Pair @@ -53,6 +55,7 @@ -- TODO: should we use @(Data.Type.Equality.:~:)@ everywhere instead of our own 'TypeEq'? import Unsafe.Coerce +import Data.Proxy import Language.Hakaru.Syntax.IClasses import Language.Hakaru.Types.DataKind ---------------------------------------------------------------- @@ -281,6 +284,11 @@ sSymbol_Maybe :: Sing "Maybe" sSymbol_Maybe = SingSymbol +someSSymbol :: String -> (forall s . Sing (s :: Symbol) -> k) -> k +someSSymbol s k = case TL.someSymbolVal s of { TL.SomeSymbol (_::Proxy s) -> k (SingSymbol :: Sing s) } + +ssymbolVal :: forall s. Sing (s :: Symbol) -> String +ssymbolVal SingSymbol = TL.symbolVal (Proxy :: Proxy s) instance Eq (Sing (s :: Symbol)) where (==) = eq1 diff --git a/haskell/System/MapleSSH.hs b/haskell/System/MapleSSH.hs --- a/haskell/System/MapleSSH.hs +++ b/haskell/System/MapleSSH.hs @@ -1,4 +1,4 @@ -module System.MapleSSH (maple) where +module System.MapleSSH (maple, mapleWithArgs) where import Data.Maybe(fromMaybe) import Data.Char(isSpace) @@ -25,20 +25,22 @@ return (ssh, user, server, command) where get name def = fmap (fromMaybe def) (lookupEnv name) -process :: IO CreateProcess -process = do +processWithArgs :: [String] -> IO CreateProcess +processWithArgs args = do bin <- lookupEnv "LOCAL_MAPLE" case bin of - Just b -> return $ proc b ["-q", "-t"] + Just b -> return $ proc b args Nothing -> do (ssh, user, server, command) <- envVarsSSH - let commands = command ++ " -q -t" -- quiet mode + let commands = command ++ concatMap (' ':) args return $ proc ssh ["-l" ++ user, server, commands] - maple :: String -> IO String -maple cmd = do - p <- process +maple = flip mapleWithArgs ["-q", "-t"] + +mapleWithArgs :: String -> [String] -> IO String +mapleWithArgs cmd args = do + p <- processWithArgs args (Just inH, Just outH, Nothing, p') <- createProcess p { std_in = CreatePipe, std_out = CreatePipe, close_fds = True } hPutStrLn inH $ cmd ++ ";" hClose inH diff --git a/haskell/Tests/ASTTransforms.hs b/haskell/Tests/ASTTransforms.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/ASTTransforms.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE RankNTypes #-} + +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Tests.ASTTransforms (allTests) where + +import Control.Monad +import qualified Data.Number.LogFloat as LF +import qualified Data.Vector as V +import GHC.Word (Word32) +import Language.Hakaru.Sample (runEvaluate) +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.ANF (normalize) +import Language.Hakaru.Syntax.CSE (cse) +import Language.Hakaru.Syntax.Prune (prune) +import Language.Hakaru.Syntax.Hoist (hoist) +import Language.Hakaru.Syntax.Uniquify (uniquify) +import Language.Hakaru.Syntax.Unroll (unroll) +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.AST.Eq (alphaEq) +import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.DatumCase +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Syntax.Prelude +import Language.Hakaru.Syntax.Value +import Language.Hakaru.Syntax.Variable +import Language.Hakaru.Types.Coercion +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.HClasses +import Language.Hakaru.Types.Sing +import Prelude hiding (product, (*), (+), + (-), (==)) + +import qualified System.Random.MWC as MWC +import Test.HUnit +import Tests.Disintegrate hiding (allTests) +import Tests.TestTools + +checkMeasure :: String + -> Value ('HMeasure a) + -> Value ('HMeasure a) + -> Assertion +checkMeasure p (VMeasure m1) (VMeasure m2) = do + -- Generate 2 copies of the same random seed so that sampling the random seeds + -- always produce the same trace of results. + g1 <- MWC.createSystemRandom + s <- MWC.save g1 + g2 <- MWC.restore s + forM_ [1 :: Int .. 10000] $ \_ -> do + p1 <- LF.logFloat `fmap` MWC.uniform g1 + p2 <- LF.logFloat `fmap` MWC.uniform g2 + Just (v1, w1) <- m1 (VProb p1) g1 + Just (v2, w2) <- m2 (VProb p2) g2 + assertEqual p v1 v2 + assertEqual p w1 w2 + +allTests :: Test +allTests = test [ TestLabel "ANF" anfTests ] + +opts :: (ABT Term abt) => abt '[] a -> abt '[] a +opts = uniquify . prune . cse . hoist . uniquify . normalize + +optsUnroll :: (ABT Term abt) => abt '[] a -> abt '[] a +optsUnroll = uniquify . prune . cse . normalize . unroll + +anfTests :: Test +anfTests = test [ "example1" ~: testNormalizer "example1" example1 example1' + , "example2" ~: testNormalizer "example2" example2 example2' + , "example3" ~: testNormalizer "example3" example3 example3' + + -- Test some deterministic results + , "runExample1" ~: testPreservesResult "example1" example1 normalize + , "runExample2" ~: testPreservesResult "example2" example2 normalize + , "runExample3" ~: testPreservesResult "example3" example3 normalize + + -- Test some programs which produce measures, these are + -- statistical tests + , "norm1a" ~: testPreservesMeasure "norm1a" norm1a normalize + , "norm1b" ~: testPreservesMeasure "norm1b" norm1b normalize + , "norm1c" ~: testPreservesMeasure "norm1c" norm1c normalize + , "easyRoad" ~: testPreservesMeasure "easyRoad" easyRoad normalize + , "helloWorld100" ~: testPreservesMeasure "helloWorld100" helloWorld100 normalize + + -- Test some deterministic results + , "runExample1CSE" ~: testPreservesResult "example1" example1 opts + , "runExample2CSE" ~: testPreservesResult "example2" example2 opts + , "runExample3CSE" ~: testPreservesResult "example3" example3 opts + + , "cse1" ~: testCSE "cse1" example1CSE example1CSE' + , "cse2" ~: testCSE "cse2" example2CSE example2CSE' + , "cse3" ~: testCSE "cse3" example3CSE example3CSE + , "cse4" ~: testCSE "cse4" (normalize example3CSE) example2CSE' + + -- Test some programs which produce measures, these are + -- statistical tests + , "norm1a all" ~: testPreservesMeasure "norm1a" norm1a opts + , "norm1b all" ~: testPreservesMeasure "norm1b" norm1b opts + , "norm1c all" ~: testPreservesMeasure "norm1c" norm1c opts + , "easyRoad all" ~: testPreservesMeasure "easyRoad" easyRoad opts + , "helloWorld100 all" ~: testPreservesMeasure "helloWorld100" helloWorld100 opts + + , "example1Hoist" ~: testPreservesResult "result" example1Hoist opts + , "example1Hoist" ~: testTransform "transform" example1Hoist example1Hoist' opts + + , "unroll" ~: testTransform "unroll" example1Unroll example1Unroll' optsUnroll + ] + + +example1 :: TrivialABT Term '[] 'HReal +example1 = if_ (real_ 1 == real_ 2) + (real_ 2 + real_ 3) + (real_ 3 + real_ 4) + +example1' :: TrivialABT Term '[] 'HReal +example1' = let_ (real_ 1 == real_ 2) $ \v -> + if_ v (real_ 2 + real_ 3) + (real_ 3 + real_ 4) + +example2 :: TrivialABT Term '[] 'HNat +example2 = let_ (nat_ 1) $ \ a -> triv ((summate a (a + (nat_ 10)) (\i -> i)) + + (product a (a + (nat_ 10)) (\i -> i))) + +example2' :: TrivialABT Term '[] 'HNat +example2' = let_ (nat_ 1) $ \ x4 -> + let_ (x4 + nat_ 10) $ \ x3 -> + let_ (summate x4 x3 (\ x0 -> x0)) $ \ x2 -> + let_ (x4 + nat_ 10) $ \ x1 -> + let_ (product x4 x1 (\ x0 -> x0)) $ \ x0 -> + x2 + x0 + +example3 :: TrivialABT Term '[] 'HReal +example3 = triv (real_ 1 * (real_ 2 + real_ 3) * (real_ 4 + (real_ 5 + (real_ 6 * real_ 7)))) + + +example3' :: TrivialABT Term '[] 'HReal +example3' = let_ (real_ 2 + real_ 3) $ \ x2 -> + let_ (real_ 6 * real_ 7) $ \ x1 -> + let_ (real_ 4 + real_ 5 + x1) $ \ x0 -> + real_ 1 * x2 * x0 + +testNormalizer :: (ABT Term abt) => String -> abt '[] a -> abt '[] a -> Assertion +testNormalizer name a b = testTransform name a b normalize + +testTransform + :: (ABT Term abt) + => String + -> abt '[] a + -> abt '[] a + -> (abt '[] a -> abt '[] a) + -> Assertion +testTransform name a b opt = assertBool name (alphaEq (opt a) b) + +testCSE :: (ABT Term abt) => String -> abt '[] a -> abt '[] a -> Assertion +testCSE name a b = assertBool name (alphaEq (cse a) b) + +testPreservesResult + :: forall (a :: Hakaru) abt . (ABT Term abt) + => String + -> abt '[] a + -> (abt '[] a -> abt '[] a) + -> Assertion +testPreservesResult name ast opt = assertEqual name result1 result2 + where result1 = runEvaluate ast + result2 = runEvaluate (opt ast) + +testPreservesMeasure + :: forall (a :: Hakaru) abt . (ABT Term abt) + => String + -> abt '[] ('HMeasure a) + -> (abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)) + -> Assertion +testPreservesMeasure name ast opt = checkMeasure name result1 result2 + where result1 = runEvaluate ast + result2 = runEvaluate (opt ast) + +example1CSE :: TrivialABT Term '[] 'HReal +example1CSE = let_ (real_ 1 + real_ 2) $ \x -> + let_ (real_ 1 + real_ 2) $ \y -> + x + y + +example1CSE' :: TrivialABT Term '[] 'HReal +example1CSE' = let_ (real_ 1 + real_ 2) $ \x -> + x + x + +example2CSE :: TrivialABT Term '[] 'HReal +example2CSE = let_ (summate (nat_ 0) (nat_ 1) $ \x -> real_ 1) $ \x -> + let_ (summate (nat_ 0) (nat_ 1) $ \x -> real_ 1) $ \y -> + x + y + +example2CSE' :: TrivialABT Term '[] 'HReal +example2CSE' = let_ (summate (nat_ 0) (nat_ 1) $ \x -> real_ 1) $ \x -> + x + x + +example3CSE :: TrivialABT Term '[] 'HReal +example3CSE = (summate (nat_ 0) (nat_ 1) $ \x -> real_ 1) + + (summate (nat_ 0) (nat_ 1) $ \x -> real_ 1) + + +example1Unroll :: TrivialABT Term '[] 'HInt +example1Unroll = (summate (int_ 0) (int_ 100) $ \x -> x + (int_ 1 * int_ 42)) + +example1Unroll' :: TrivialABT Term '[] 'HInt +example1Unroll' = let_ (int_ 0 == int_ 100) $ \cond -> + if_ cond (int_ 0) + (let_ (int_ 1 * int_ 42) $ \tmp -> + let_ (int_ 0 + tmp) $ \first -> + let_ (int_ 0 + int_ 1) $ \start -> + let_ (summate start (int_ 100) $ (+ tmp)) $ \total -> + first + total) + +example1Hoist :: TrivialABT Term '[] 'HInt +example1Hoist = summate (int_ 0) (int_ 1) $ \_ -> + summate (int_ 1) (int_ 2) id + +example1Hoist' :: TrivialABT Term '[] 'HInt +example1Hoist' = let_ (summate (int_ 1) (int_ 2) id) $ \x -> + summate (int_ 0) (int_ 1) (const x) + diff --git a/haskell/Tests/Disintegrate.hs b/haskell/Tests/Disintegrate.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/Disintegrate.hs @@ -0,0 +1,571 @@ +{-# LANGUAGE DataKinds + , GADTs + , TypeOperators + , NoImplicitPrelude + , FlexibleContexts + , OverloadedStrings + #-} + +module Tests.Disintegrate where + +import Prelude (($), (.), (++), head, String, Maybe(..)) +import qualified Prelude +import qualified Data.List.NonEmpty as L + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.Datum +import Language.Hakaru.Syntax.Prelude +import Language.Hakaru.Syntax.IClasses (Some2(..), TypeEq(..), jmEq1) +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.Sing +import Language.Hakaru.Pretty.Concrete +import Language.Hakaru.Syntax.TypeCheck +import Language.Hakaru.Evaluation.Types (fromWhnf) +import Language.Hakaru.Evaluation.DisintegrationMonad (runDis) +import Language.Hakaru.Disintegrate + +import Test.HUnit +import Tests.TestTools +import Tests.Models hiding (easyRoad) + +---------------------------------------------------------------- +type Model a b = TrivialABT Term '[] ('HMeasure (HPair a b)) +type Cond a b = TrivialABT Term '[] (a ':-> 'HMeasure b) +---------------------------------------------------------------- + +-- | A very simple program. Is sufficient for testing escape and +-- capture of substitution. +norm0a :: Model 'HReal 'HReal +norm0a = + normal (real_ 0) (prob_ 1) >>= \x -> + normal x (prob_ 1) >>= \y -> + dirac (pair y x) + +-- | A version of 'norm0' which adds a type annotation at the +-- top-level; useful for testing that using 'Ann_' doesn't cause +-- perform\/disintegrate to loop. +norm0b :: Model 'HReal 'HReal +norm0b = ann_ sing norm0a + +-- | A version of 'norm0' which inserts an annotation around the +-- 'Datum' constructor itself. The goal here is to circumvent the +-- @typeOf_{Datum_}@ issue without needing to change the 'Datum' +-- type nor the 'typeOf' definition. +norm0c :: Model 'HReal 'HReal +norm0c = + normal (real_ 0) (prob_ 1) >>= \x -> + normal x (prob_ 1) >>= \y -> + dirac (ann_ sing $ pair y x) + +-- | What we expect 'norm0a' (and variants) to disintegrate to. +norm0' :: Cond 'HReal 'HReal +norm0' = + lam $ \y -> + normal (real_ 0) (prob_ 1) >>= \x -> + weight (densityNormal x (prob_ 1) y) >> + dirac x + +{- +-- Eliminating some redexes of 'norm0'', that is: + lam $ \y -> + normal (real_ 0) (prob_ 1) >>= \x -> + withWeight + (exp ((x - y) ^ nat_ 2 / real_ 2) + / (nat_ 2 `thRootOf` (prob_ 2 * pi))) + $ dirac x + +-- N.B., calling 'evaluate' on 'norm0'' doesn't catch those redexes because they're not on the way to computing stuff. need to call 'constantPropagation' to get rid of them. +-} + + +testPerform0a, testPerform0b, testPerform0c + :: [Model 'HReal 'HReal] +testPerform0a = runPerform norm0a +testPerform0b = runPerform norm0b +testPerform0c = runPerform norm0c + +testDisintegrate0a, testDisintegrate0b, testDisintegrate0c + :: [Cond 'HReal 'HReal] +testDisintegrate0a = disintegrate norm0a +testDisintegrate0b = disintegrate norm0b +testDisintegrate0c = disintegrate norm0c + +-- | The goal of this test is to be sure we maintain proper hygiene +-- for the weight component when disintegrating superpose. Moreover, +-- in earlier versions it used to throw VarEqTypeError due to +-- 'disintegrate' not choosing a sufficiently fresh variable name +-- for its lambda; thus this also serves as a regression test to +-- make sure we don't run into that problem again. +testHygiene0b :: [Cond 'HReal 'HReal] +testHygiene0b = + disintegrate $ + let_ (prob_ 1) $ \x -> + withWeight x norm0b + +---------------------------------------------------------------- +-- | This simple progam is to check for disintegrating case analysis +-- when the scrutinee contains a heap-bound variable. +norm1a :: Model 'HReal HUnit +norm1a = + normal (real_ 3) (prob_ 2) >>= \x -> + dirac $ if_ (x < real_ 0) + (ann_ sing $ pair (negate x) unit) + (ann_ sing $ pair x unit) + +norm1b :: Model HReal HUnit +norm1b = + normal (real_ 3) (prob_ 2) >>= \x -> + if_ (x < real_ 0) + (ann_ sing . dirac $ pair (negate x) unit) + (ann_ sing . dirac $ pair x unit) + +norm1c :: Model 'HReal HUnit +norm1c = + normal (real_ 3) (prob_ 2) >>= \x -> + if_ (x < real_ 0) + (dirac . ann_ sing $ pair (negate x) unit) + (dirac . ann_ sing $ pair x unit) + +norm1' :: Cond 'HReal HUnit +norm1' = + lam $ \t -> superpose $ + L.fromList + [ (one, + weight (densityNormal (real_ 3) (prob_ 2) (negate t)) >>= \_ -> + let_ (negate t) $ \x -> + case_ (x < zero) [branch pTrue (dirac unit)]) + , (one, + weight (densityNormal (real_ 3) (prob_ 2) t) >>= \_ -> + case_ (t < zero) [branch pFalse (dirac unit)]) ] + +-- BUG: the first solutions returned by 'testPerform1b' and 'testPerform1c' break hygiene! They drops the variable bound by 'normal' and has all the uses of @x@ become free. +testPerform1a, testPerform1b, testPerform1c + :: [Model 'HReal HUnit] +testPerform1a = runPerform norm1a +testPerform1b = runPerform norm1b +testPerform1c = runPerform norm1c + +testDisintegrate1a, testDisintegrate1b, testDisintegrate1c + :: [Cond 'HReal HUnit] +testDisintegrate1a = disintegrate norm1a +testDisintegrate1b = disintegrate norm1b +testDisintegrate1c = disintegrate norm1c + + +---------------------------------------------------------------- +norm2 :: Model 'HReal 'HReal +norm2 = + normal (real_ 3) (prob_ 2) >>= \x -> + normal (real_ 5) (prob_ 4) >>= \y -> + dirac $ if_ (x < y) + (pair y x) + (pair x x) + +norm2' :: Cond 'HReal 'HReal +norm2' = + lam $ \t -> superpose $ + L.fromList + [ (one, + normal (real_ 3) (prob_ 2) >>= \x -> + weight (densityNormal (real_ 5) (prob_ 4) t) >>= \_ -> + case_ (x < t) [branch pTrue (dirac x)]) + , (one, + weight (densityNormal (real_ 3) (prob_ 2) t) >>= \_ -> + normal (real_ 5) (prob_ 4) >>= \y -> + case_ (t < y) [branch pFalse (dirac t)]) ] + +testPerform2 + :: [Model 'HReal 'HReal] +testPerform2 = runPerform norm2 + +testDisintegrate2 + :: [Cond 'HReal 'HReal] +testDisintegrate2 = disintegrate norm2 + +---------------------------------------------------------------- + +normSquare :: Model 'HProb 'HReal +normSquare = + normal (real_ 0) (prob_ 1) >>= \x -> + dirac (pair (square x) x) + +normSquare' :: Cond 'HProb 'HReal +normSquare' = + lam $ \t -> + weight (recip (nat2prob (nat_ 2) * sqrt t)) >>= \_ -> + weight (densityNormal (real_ 0) (prob_ 1) (fromProb (sqrt t))) >>= \_ -> + dirac (fromProb (sqrt t)) >>= \x -> + dirac x + +---------------------------------------------------------------- + +normDirac :: Model 'HReal 'HReal +normDirac = + normal (real_ 0) (prob_ 1) >>= \x -> + dirac x >>= \y -> + dirac (pair y x) + +normDirac' :: Cond 'HReal 'HReal +normDirac' = + lam $ \t -> + weight (densityNormal (real_ 0) (prob_ 1) t) >>= \_ -> + dirac t >>= \x -> + dirac x + +---------------------------------------------------------------- + +pendulum :: Model 'HReal 'HReal +pendulum = + normal (real_ 42) (prob_ 1) >>= \theta -> + dirac (sin theta) >>= \x -> + normal (real_ 0) (prob_ 1) >>= \noise -> + dirac (pair (x + noise) theta) + +pendulum' :: Cond 'HReal 'HReal +pendulum' = + lam $ \t -> + normal (real_ 42) (prob_ 1) >>= \theta -> + weight (densityNormal (real_ 0) (prob_ 1) (t - sin theta)) >>= \_ -> + dirac theta + +---------------------------------------------------------------- +easyRoad :: Model (HPair 'HReal 'HReal) (HPair 'HProb 'HProb) +easyRoad = + uniform (real_ 3) (real_ 8) >>= \noiseT_ -> + uniform (real_ 1) (real_ 4) >>= \noiseE_ -> + let_ (unsafeProb noiseT_) $ \noiseT -> + let_ (unsafeProb noiseE_) $ \noiseE -> + normal (real_ 0) noiseT >>= \x1 -> + normal x1 noiseE >>= \m1 -> + normal x1 noiseT >>= \x2 -> + normal x2 noiseE >>= \m2 -> + dirac (pair (pair m1 m2) (pair noiseT noiseE)) + +easyRoad' :: Cond (HPair 'HReal 'HReal) (HPair 'HProb 'HProb) +easyRoad' = + lam $ \t -> + unpair t (\t1 t2 -> + uniform (real_ 3) (real_ 8) >>= \noiseT_ -> + uniform (real_ 1) (real_ 4) >>= \noiseE_ -> + let_ (unsafeProb noiseT_) $ \noiseT -> + let_ (unsafeProb noiseE_) $ \noiseE -> + normal (real_ 0) noiseT >>= \x1 -> + weight (densityNormal x1 noiseE t1) >>= \_ -> + normal x1 noiseT >>= \x2 -> + weight (densityNormal x2 noiseE t2) >> + dirac (pair noiseT noiseE)) + +testPerformEasyRoad + :: [Model (HPair 'HReal 'HReal) (HPair 'HProb 'HProb)] +testPerformEasyRoad = runPerform easyRoad + + +testDisintegrateEasyRoad + :: [Cond (HPair 'HReal 'HReal) (HPair 'HProb 'HProb)] +testDisintegrateEasyRoad = disintegrate easyRoad + +---------------------------------------------------------------- +helloWorld100 + :: Model ('HArray 'HReal) 'HReal +helloWorld100 = + normal (real_ 0) (prob_ 1) >>= \mu -> + plate (nat_ 100) (\_ -> normal mu (prob_ 1)) >>= \v -> + dirac (pair v mu) + +helloWorld100' + :: Cond ('HArray 'HReal) 'HReal +helloWorld100' = + lam $ \t -> + normal (real_ 0) (prob_ 1) >>= \mu -> + plate (nat_ 100) + (\i -> weight (densityNormal mu (prob_ 1) (t ! i))) >> + dirac mu + +testHelloWorld100 + :: [Cond ('HArray 'HReal) 'HReal] +testHelloWorld100 = disintegrate helloWorld100 + +---------------------------------------------------------------- +copy1 :: Model ('HArray 'HReal) HUnit +copy1 = + plate n (\_ -> normal (real_ 0) (prob_ 1)) >>= \u -> + dirac (array n (\i -> u ! i)) >>= \v -> + dirac (pair v unit) + where n = nat_ 100 + +copy1' :: Cond ('HArray 'HReal) HUnit +copy1' = + lam $ \t -> + plate (nat_ 100) + (\i -> weight (densityNormal (real_ 0) (prob_ 1) (t ! i))) >> + dirac unit + +testCopy1 :: [Cond ('HArray 'HReal) HUnit] +testCopy1 = disintegrate copy1 + +---------------------------------------------------------------- +copy2 :: Model ('HArray 'HReal) HUnit +copy2 = + plate n (\_ -> normal (real_ 0) (prob_ 1)) >>= \u -> + plate n (\j -> dirac (u ! j)) >>= \v -> + dirac (pair v unit) + where n = nat_ 100 + +testCopy2 :: [Cond ('HArray 'HReal) HUnit] +testCopy2 = disintegrate copy2 + + +---------------------------------------------------------------- +naiveBayes + :: Model ('HArray ('HArray 'HNat)) ('HArray 'HNat) +naiveBayes = + plate numLabels (\_ -> dirichlet (array sizeVocab (\_ -> prob_ 1))) >>= \bs -> + dirichlet (array numLabels (\_ -> prob_ 1)) >>= \ts -> + plate numDocs (\_ -> categorical ts) >>= \zs -> + plate numDocs (\i -> plate sizeEachDoc + (\_ -> categorical (bs ! (zs ! i)))) >>= \ds -> + dirac (pair ds zs) + where sizeVocab = nat_ 1000 + numLabels = nat_ 40 + numDocs = nat_ 200 + sizeEachDoc = nat_ 5000 + +naiveBayes' + :: Cond ('HArray ('HArray 'HNat)) ('HArray 'HNat) +naiveBayes' = + lam $ \t -> + Prelude.error "TODO define naiveBayes'" + +testNaiveBayes + :: [Cond ('HArray ('HArray 'HNat)) ('HArray 'HNat)] +testNaiveBayes = disintegrate naiveBayes + +---------------------------------------------------------------- +-- | R2 benchmarks +-- Based on examples from the R2 probabilistic programming tool +-- Found in r2-0.0.1/examples/ when downloaded from: +-- https://www.microsoft.com/en-us/download/details.aspx?id=52372 + +clinicalTrial + :: Model (HPair ('HArray HBool) ('HArray HBool)) HBool +clinicalTrial = + bern (prob_ 0.5) >>= \isEffective -> + beta (prob_ 1) (prob_ 1) >>= \probControl -> + beta (prob_ 1) (prob_ 1) >>= \probTreated -> + beta (prob_ 1) (prob_ 1) >>= \probAll -> + if_ isEffective + (liftM2 pair (plate n (\_ -> bern probControl)) + (plate m (\_ -> bern probTreated))) + (liftM2 pair (plate n (\_ -> bern probAll)) + (plate m (\_ -> bern probAll))) >>= \groups -> + dirac (pair groups isEffective) + where (n,m) = (nat_ 1000, nat_ 1000) + +coinBias :: Model ('HArray HBool) 'HProb +coinBias = beta (prob_ 2) (prob_ 5) >>= \bias -> + plate (nat_ 5) (\_ -> bern bias) >>= \tossResults -> + dirac (pair tossResults bias) + +digitRecognition :: Model ('HArray HBool) 'HNat +digitRecognition = + categorical dataPrior >>= \y -> + plate n (\i -> bern $ (dataParams ! y) ! i) >>= \x -> + dirac (pair x y) + where n = nat_ 784 + dataPrior = var (Variable "dataPrior" 73 (SArray SProb)) + dataParams = var (Variable "dataParams" 41 (SArray (SArray SProb))) + +hiv :: Model ('HArray 'HReal) + (HPair (HPair ('HArray 'HReal) ('HArray 'HReal)) + ('HArray 'HReal)) +hiv = normal (real_ 0) (prob_ 1) >>= \muA1 -> + normal (real_ 0) (prob_ 1) >>= \muA2 -> + uniform (real_ 0) (real_ 100) >>= \sigmaA1 -> + uniform (real_ 0) (real_ 100) >>= \sigmaA2 -> + plate (nat_ 84) (\_ -> normal muA1 (unsafeProb sigmaA1)) >>= \a1 -> + plate (nat_ 84) (\_ -> normal ((real_ 0.1)*muA2) (unsafeProb sigmaA2)) >>= \a2 -> + dirac (array n (\i -> a1 ! (unsafeMinusNat (dataPerson ! i) (nat_ 1)) + + a2 ! (unsafeMinusNat (dataPerson ! i) (nat_ 1)) * + dataTime ! i)) >>= \yHat -> + uniform (real_ 0) (real_ 100) >>= \sigmaY -> + plate n (\i -> normal (yHat ! i) (unsafeProb sigmaY)) >>= \y -> + dirac (pair y (pair (pair a1 a2) + (arrayLit [muA1, muA2, sigmaA1, sigmaA2, sigmaY]))) + where n = nat_ 369 + dataPerson = var (Variable "dataPerson" 73 (SArray SNat)) + dataTime = var (Variable "dataTime" 41 (SArray SReal)) -- hacks :( + +linearRegression + :: Model ('HArray 'HReal) ('HArray 'HReal) +linearRegression = + normal (real_ 0) (prob_ 1) >>= \a -> + normal (real_ 5) (prob_ 1.82574185835055371152) >>= \b -> + gamma (prob_ 1) (prob_ 1) >>= \invNoise -> + plate n (\i -> normal (a * (dataX ! i) + b) (recip $ sqrt invNoise)) >>= \y -> + dirac (pair y (arrayLit [a, b, fromProb invNoise])) + where n = nat_ 1000 + dataX = var (Variable "dataX" 73 (SArray SReal)) -- hack :( + +surveyUnbias :: Model ('HArray HBool) + (HPair ('HArray ('HArray 'HProb)) ('HArray 'HReal)) +surveyUnbias = + dirac (size population) >>= \n -> + plate n (\_ -> beta (prob_ 1) (prob_ 1)) >>= \bias -> + dirac (array n (\i -> population!i * bias!i)) >>= \mean -> + dirac (array n (\i -> mean!i `unsafeMinusProb` + (mean!i * bias!i))) >>= \variance -> + plate n (\i -> normal (fromProb $ mean!i) + (sqrt $ variance!i)) >>= \votes -> + dirac (size personGender) >>= \m -> + dirac (array m (\i -> bias ! (personGender ! i))) >>= \ansBias -> + plate m (\i -> bern $ ansBias ! i) >>= \answer -> + dirac (pair answer (pair (arrayLit [bias, mean, variance]) votes)) + where population = var (Variable "population" 73 (SArray SProb)) + personGender = var (Variable "personGender" 41 (SArray SNat)) + +surveyUnbias2 :: Model ('HArray HBool) + (HPair ('HArray 'HProb) ('HArray 'HInt)) +surveyUnbias2 = + dirac (size population) >>= \n -> + plate n (\_ -> beta (prob_ 1) (prob_ 1)) >>= \bias -> + plate n (\i -> binomial (population!i) (bias!i)) >>= \votes -> + dirac (size personGender) >>= \m -> + dirac (array m (\i -> bias ! (personGender ! i))) >>= \ansBias -> + plate m (\i -> bern $ ansBias ! i) >>= \answer -> + dirac (pair answer (pair bias votes)) + where population = var (Variable "population" 73 (SArray SNat)) + personGender = var (Variable "personGender" 41 (SArray SNat)) + +---------------------------------------------------------------- + +unzipFst :: Model ('HArray 'HReal) HUnit +unzipFst = plate n (\_ -> liftM2 pair (normal zero one) + (normal zero one)) >>= \u -> + dirac (array n (\i -> fst (u ! i))) >>= \v -> + dirac (pair v unit) + where n = nat_ 1000 + +transpose :: Model ('HArray ('HArray 'HReal)) HUnit +transpose = plate n (\_ -> plate n (\_ -> normal zero one)) >>= \u -> + dirac (array n (\i -> array n (\j -> (u ! j) ! i))) >>= \v -> + dirac (pair v unit) + where n = nat_ 3500 + +---------------------------------------------------------------- + +testEmissions :: Model ('HArray 'HReal) HUnit +testEmissions = plate n (\_ -> lebesgue) >>= \xs -> + plate n (\_ -> lebesgue) >>= \ys -> + dirac (pair (array n (\i -> (xs ! i) + (ys ! i))) unit) + where n = nat_ 100 + +-- Tug of war example for debugging scope extrusion errors +-- https://github.com/hakaru-dev/hakaru/issues/30 +tow :: Model (HPair 'HReal 'HReal) 'HReal +tow = normal zero one >>= \alice -> + normal zero one >>= \bob -> + normal zero one >>= \carol -> + (normal alice one >>= \a -> + normal bob one >>= \b -> + dirac (a-b)) >>= \match1 -> + (normal bob one >>= \b -> + normal carol one >>= \c -> + dirac (b-c)) >>= \match2 -> + (normal alice one >>= \a -> + normal carol one >>= \c -> + dirac (a-c)) >>= \match3 -> + dirac (pair (pair match1 match2) match3) + +-- Smaller version of tug of war +minimaltow :: Model 'HReal HUnit +minimaltow = normal zero one >>= \alice -> + normal zero one >>= \bob -> + (normal alice one >>= \a -> + normal bob one >>= \b -> + dirac (a-b)) >>= \match -> + dirac (pair match unit) + +slice :: Model 'HReal 'HReal +slice = normal zero one >>= \x -> + uniform zero (fromProb (densityNormal zero one x)) >>= \y -> + dirac (pair y x) + +oneAndAll :: Model 'HReal ('HArray 'HReal) +oneAndAll = plate (nat_ 100) (\_ -> normal zero one) >>= \x -> + dirac (pair (x ! nat_ 3) x) + +runPerform + :: TrivialABT Term '[] ('HMeasure a) + -> [TrivialABT Term '[] ('HMeasure a)] +runPerform e = runDis (fromWhnf `Prelude.fmap` perform e) [Some2 e] + +-- | Tests that disintegration doesn't error and produces at least +-- one solution. +testDis + :: (ABT Term abt) + => String + -> abt '[] ('HMeasure (HPair a b)) + -> Assertion +testDis p = + assertBool (p ++ ": no disintegration found") + . Prelude.not + . Prelude.null + . disintegrate + +showFirst :: Model a b -> Prelude.IO () +showFirst e = let anss = disintegrate e + in if Prelude.null anss + then Prelude.putStrLn $ "no disintegration found" + else Prelude.print $ pretty (head anss) + +-- TODO: put all the "perform" tests in here +allTests :: Test +allTests = test + [ testDis "testDisintegrate0a" norm0a + , testDis "testDisintegrate0b" norm0b + , testDis "testDisintegrate0c" norm0c + , assertAlphaEq "testDisintegrate0a" (head testDisintegrate0a) norm0' + , assertAlphaEq "testDisintegrate0b" (head testDisintegrate0b) norm0' + , assertAlphaEq "testDisintegrate0c" (head testDisintegrate0c) norm0' + , assertBool "testHygiene0b" $ Prelude.not (Prelude.null testHygiene0b) + , testDis "testDisintegrate1a" norm1a + , testDis "testDisintegrate1b" norm1b + , testDis "testDisintegrate1c" norm1c + , assertAlphaEq "testDisintegrate1a" (head testDisintegrate1a) norm1' + , assertAlphaEq "testDisintegrate1b" (head testDisintegrate1b) norm1' + , assertAlphaEq "testDisintegrate1c" (head testDisintegrate1c) norm1' + , testDis "testDisintegrate2" norm2 + , assertAlphaEq "testDisintegrate2" (head testDisintegrate2) norm2' + , testWithConcrete' match_norm_unif LaxMode $ \_typ ast -> + case jmEq1 _typ (SMeasure $ sPair SReal sBool) of + Just Refl -> testDis "testMatchNormUnif" ast + Nothing -> assertFailure "BUG: jmEq1 got the wrong type" + , testWithConcrete' dont_atomize_weights LaxMode $ \_typ ast -> + case jmEq1 _typ (SMeasure $ sPair SReal sUnit) of + Just Refl -> testDis "testAtomizeWeights" ast + Nothing -> assertFailure "BUG: jmEq1 got the wrong type" + , assertBool "testPerformEasyRoad" $ Prelude.not (Prelude.null testPerformEasyRoad) + , testDis "testDisintegrateEasyRoad" easyRoad + , assertAlphaEq "testDisintegrateEasyRoad" (head testDisintegrateEasyRoad) easyRoad' + , testDis "testHelloWorld100" helloWorld100 + , assertAlphaEq "testHelloWorld100" (head testHelloWorld100) helloWorld100' + , testDis "testCopy1" copy1 + , assertAlphaEq "testCopy1" (head testCopy1) copy1' + , testDis "testCopy2" copy2 + , assertAlphaEq "testCopy2" (head testCopy2) copy1' + , testDis "testPendulum" pendulum + , assertAlphaEq "testPendulumDisintegrate" (disintegrate pendulum Prelude.!! 2) pendulum' + , testDis "testNaiveBayes" naiveBayes + , testDis "testClinicalTrial" clinicalTrial + , testDis "testCoinBias" coinBias + , testDis "testDigitRecognition" digitRecognition + , testDis "TestHIV" hiv + , testDis "testLinearRegression" linearRegression + , testDis "testSurveyUnbias" surveyUnbias + , testDis "testSurveyUnbias2" surveyUnbias2 + ] + + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/haskell/Tests/Models.hs b/haskell/Tests/Models.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/Models.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE DataKinds + , NoImplicitPrelude + , OverloadedStrings + , FlexibleContexts #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +---------------------------------------------------------------- +-- 2016.04.22 +-- | +-- Module : Tests.Models +-- Copyright : Copyright (c) 2016 the Hakaru team +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : experimental +-- Portability : GHC-only +-- +-- Some common models used in many different tests. +-- +-- TODO: we might should adjust our use of 'pair' to include a type +-- annotation, to avoid issues about 'Language.Hakaru.Syntax.TypeOf.typeOf' +-- on 'Datum'. +---------------------------------------------------------------- +module Tests.Models where + +import Data.Text +import qualified Data.List.NonEmpty as L + +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Syntax.AST (Term) +import Language.Hakaru.Syntax.ABT (ABT) +import Language.Hakaru.Syntax.Prelude + + +---------------------------------------------------------------- +---------------------------------------------------------------- +uniform_0_1 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +uniform_0_1 = uniform zero one + +-- build uniform with nats and coercions +uniformC :: (ABT Term abt) + => abt '[] 'HNat + -> abt '[] 'HNat + -> abt '[] ('HMeasure 'HReal) +uniformC lo hi = uniform (nat2real lo) (nat2real hi) + +normal_0_1 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +normal_0_1 = normal zero one + +-- build normal with nats and coercions +normalC :: (ABT Term abt) + => abt '[] 'HNat + -> abt '[] 'HNat + -> abt '[] ('HMeasure 'HReal) +normalC mu sd = normal (nat2real mu) (nat2prob sd) + +gamma_1_1 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +gamma_1_1 = gamma one one + +beta_1_1 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +beta_1_1 = beta one one + + +-- TODO: a better name for this +t4 :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HProb HBool)) +t4 = + beta one one >>= \bias -> + bern bias >>= \coin -> + dirac (pair bias coin) + +t4' :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HProb HBool)) +t4' = + uniform zero one >>= \x -> + let x' = unsafeProb x in + superpose (L.fromList + [ (x' , dirac (pair x' true)) + , (unsafeProb (one - x), dirac (pair x' false)) + ]) + +norm :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HReal 'HReal)) +norm = + normal zero one >>= \x -> + normal x one >>= \y -> + dirac (pair x y) + +unif2 :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HReal 'HReal)) +unif2 = + uniform (negate one) one >>= \x -> + uniform (x - one) (x + one) >>= \y -> + dirac (pair x y) + +match_norm_unif :: Text +match_norm_unif = unlines + [ "def bern(p prob):" + , " x <~ categorical([p, real2prob(1 - p)])" + , " return [true, false][x]" + , "" + , "x <~ bern(0.5)" + , "y <~ match x:" + , " true: normal(0,1)" + , " false: uniform(0,1)" + , "return ((y,x). pair(real, bool))" + ] + +match_norm_bool :: Text +match_norm_bool = unlines + [ "x <~ normal(3,2)" + , "return (match x < 0:" + , " true: (-x, ())" + , " false: ( x, ()))" + ] + +easyRoad :: Text +easyRoad = unlines + ["noiseT_ <~ uniform(3, 8)" + ,"noiseE_ <~ uniform(1, 4)" + ,"noiseT = unsafeProb(noiseT_)" + ,"noiseE = unsafeProb(noiseE_)" + ,"x1 <~ normal(0, noiseT)" + ,"m1 <~ normal(x1, noiseE)" + ,"x2 <~ normal(x1, noiseT)" + ,"m2 <~ normal(x2, noiseE)" + ,"return ((m1, m2), (noiseT, noiseE))" + ] + +plate_norm :: Text +plate_norm = unlines + [ "x <~ normal(0,1)" + , "y <~ plate _ of 5:" + , " normal(x,1)" + , "return (y, x)" + ] + +negate_prob :: Text +negate_prob = "unsafeProb(1.0 + negate(2.0))" + +dont_atomize_weights :: Text +dont_atomize_weights = unlines + ["x <~ uniform(1,3)" + ,"weight(real2prob(x), return (x, ()))"] + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/haskell/Tests/Parser.hs b/haskell/Tests/Parser.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/Parser.hs @@ -0,0 +1,537 @@ +{-# LANGUAGE CPP, OverloadedStrings, DataKinds #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-orphans #-} + +module Tests.Parser where + +import Prelude hiding (unlines) + +import Language.Hakaru.Parser.Parser +import Language.Hakaru.Parser.AST + +import Data.Text +import Test.HUnit +import Test.QuickCheck.Arbitrary +import Test.QuickCheck + +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative (Applicative(..), (<$>)) +#endif + +arbNat :: Gen (Positive Integer) +arbNat = arbitrary + +arbProb :: Gen (Positive Rational) +arbProb = arbitrary + +instance Arbitrary Text where + arbitrary = pack <$> ("x" ++) . show <$> getPositive <$> arbNat + shrink xs = pack <$> shrink (unpack xs) + +instance Arbitrary Literal' where + arbitrary = oneof + [ Nat <$> getPositive <$> arbNat + , Int <$> arbitrary + , Prob <$> getPositive <$> arbProb + , Real <$> arbitrary + ] + +instance Arbitrary TypeAST' where + arbitrary = frequency + [ (20, TypeVar <$> arbitrary) + , ( 1, TypeApp <$> arbitrary <*> arbitrary) + , ( 1, TypeFun <$> arbitrary <*> arbitrary) + ] + +instance Arbitrary NaryOp where + arbitrary = elements + [And, Or, Xor, Iff, Min, Max, Sum, Prod] + +instance Arbitrary a => Arbitrary (Pattern' a) where + arbitrary = oneof + [ PVar' <$> arbitrary + , return PWild' + , PData' <$> (DV <$> arbitrary <*> arbitrary) + ] + +instance Arbitrary a => Arbitrary (Branch' a) where + arbitrary = Branch' <$> arbitrary <*> arbitrary + +instance Arbitrary a => Arbitrary (AST' a) where + arbitrary = frequency + [ (10, Var <$> arbitrary) + , ( 1, Lam <$> arbitrary <*> arbitrary <*> arbitrary) + , ( 1, App <$> arbitrary <*> arbitrary) + , ( 1, Let <$> arbitrary <*> arbitrary <*> arbitrary) + , ( 1, If <$> arbitrary <*> arbitrary <*> arbitrary) + , ( 1, Ann <$> arbitrary <*> arbitrary) + , ( 1, return Infinity') + , ( 1, ULiteral <$> arbitrary) + --, ( 1, NaryOp <$> arbitrary) + , ( 1, return Empty) + , ( 1, Case <$> arbitrary <*> arbitrary) + , ( 1, Dirac <$> arbitrary) + , ( 1, Bind <$> arbitrary <*> arbitrary <*> arbitrary) + ] + +testParse :: Text -> AST' Text -> Assertion +testParse s p = + case parseHakaru s of + Left m -> assertFailure (unpack s ++ "\n" ++ show m) + Right p' -> assertEqual "" p p' + +if1, if2, if3, if4, if5 :: Text + +ifAST1, ifAST2 :: AST' Text + +if1 = "if True: 1 else: 2" + +if2 = unlines + ["if True: 1 else:" + ,"2" + ] + +if3 = unlines + ["if True:" + ," 1" + ,"else:" + ," 2" + ] + +if4 = unlines + ["if True:" + ," 1 else: 2" + ] + +if5 = unlines + ["if True:" + ," 4" + ,"else:" + ," if False:" + ," 2" + ," else:" + ," 3" + ] + +ifAST1 = + If (Var "True") + (ULiteral (Nat 1)) + (ULiteral (Nat 2)) + +ifAST2 = + If (Var "True") + (ULiteral (Nat 4)) + (If (Var "False") + (ULiteral (Nat 2)) + (ULiteral (Nat 3))) + +testIfs :: Test +testIfs = test + [ testParse if1 ifAST1 + , testParse if2 ifAST1 + , testParse if3 ifAST1 + , testParse if4 ifAST1 + , testParse if5 ifAST2 + ] + +lam1 :: Text +lam1 = "fn x nat: x+3" + +lam1AST :: AST' Text +lam1AST = Lam "x" (TypeVar "nat") + (NaryOp Sum [ Var "x" + , ULiteral (Nat 3) + ]) + +def1 :: Text +def1 = unlines + ["def foo(x nat):" + ," x + 3" + ,"foo(5)" + ] + +def2 :: Text +def2 = unlines + ["def foo(x nat): x + 3" + ,"foo(5)" + ] + +def3 :: Text +def3 = unlines + ["def foo(x real):" + ," y <~ normal(x,1.0)" + ," return (y + y. real)" + ,"foo(-2.0)" + ] + +def4 :: Text +def4 = unlines + ["def foo(x nat) nat:" + ," x+3" + ,"foo(5)" + ] + +def1AST :: AST' Text +def1AST = + Let "foo" (Lam "x" (TypeVar "nat") + (NaryOp Sum [Var "x", ULiteral (Nat 3)])) + (App (Var "foo") (ULiteral (Nat 5))) + +def2AST :: AST' Text +def2AST = + Let "foo" (Lam "x" (TypeVar "real") + (Bind "y" (App (App (Var "normal") (Var "x")) (ULiteral (Prob 1.0))) + (Dirac (Ann (NaryOp Sum [Var "y", Var "y"]) + (TypeVar "real"))))) + (App (Var "foo") (App (Var "negate") (ULiteral (Prob 2.0)))) + +def3AST :: AST' Text +def3AST = + Let "foo" (Ann + (Lam "x" (TypeVar "nat") + (NaryOp Sum [Var "x", ULiteral (Nat 3)])) + (TypeFun (TypeVar "nat") (TypeVar "nat"))) + (App (Var "foo") (ULiteral (Nat 5))) + +testLams :: Test +testLams = test + [ testParse lam1 lam1AST + , testParse def1 def1AST + , testParse def2 def1AST + , testParse def3 def2AST + , testParse def4 def3AST + ] + +let1 :: Text +let1 = unlines + ["x = 3" + ,"y = 2" + ,"x + y" + ] + +let1AST :: AST' Text +let1AST = + Let "x" (ULiteral (Nat 3)) + (Let "y" (ULiteral (Nat 2)) + (NaryOp Sum [Var "x", Var "y"])) + +testLets :: Test +testLets = test + [ testParse let1 let1AST ] + +bind1 :: Text +bind1 = unlines + ["x <~ uniform(0,1)" + ,"y <~ normal(x, 1)" + ,"dirac(y)" + ] + +bind2 :: Text +bind2 = unlines + ["x <~ uniform(0,1)" + ,"y <~ normal(x, 1)" + ,"return y" + ] + +bind1AST :: AST' Text +bind1AST = + Bind "x" (App (App (Var "uniform") + (ULiteral (Nat 0))) + (ULiteral (Nat 1))) + (Bind "y" (App (App (Var "normal") + (Var "x")) + (ULiteral (Nat 1))) + (Dirac (Var "y"))) + +ret1 :: Text +ret1 = "return return 3" + +ret1AST :: AST' Text +ret1AST = Dirac (Dirac (ULiteral (Nat 3))) + +testBinds :: Test +testBinds = test + [ testParse bind1 bind1AST + , testParse bind2 bind1AST + , testParse ret1 ret1AST + ] + +match1 :: Text +match1 = unlines + ["match e:" + ," left(a): e1" + ] + +match1AST :: AST' Text +match1AST = + Case (Var "e") + [Branch' (PData' (DV "left" [PVar' "a"])) (Var "e1")] + +-- The space between _ and : is important +match2 :: Text +match2 = unlines + ["match e:" + ," _: e" + ] + +match2AST :: AST' Text +match2AST = + Case (Var "e") + [Branch' PWild' (Var "e")] + +match3 :: Text +match3 = unlines + ["match e:" + ," a: e" + ] + +match3AST :: AST' Text +match3AST = + Case (Var "e") + [Branch' (PVar' "a") (Var "e")] + +match4 :: Text +match4 = unlines + ["match e:" + ," left(a): e1" + ," right(b): e2" + ] + +match4AST :: AST' Text +match4AST = + Case (Var "e") + [ Branch' (PData' (DV "left" [PVar' "a"])) (Var "e1") + , Branch' (PData' (DV "right" [PVar' "b"])) (Var "e2") + ] + +match5 :: Text +match5 = unlines + ["match e:" + ," left(a):" + ," e1" + ," right(b):" + ," e2" + ] + +match5AST :: AST' Text +match5AST = + Case (Var "e") + [ Branch' (PData' (DV "left" [PVar' "a"])) (Var "e1") + , Branch' (PData' (DV "right" [PVar' "b"])) (Var "e2") + ] + +match6 :: Text +match6 = unlines + ["match (2,3). pair(nat,nat):" + ," (a,b): a+b" + ] + +match6AST :: AST' Text +match6AST = + Case (Ann + (Pair + (ULiteral (Nat 2)) + (ULiteral (Nat 3))) + (TypeApp "pair" [TypeVar "nat",TypeVar "nat"])) + [Branch' (PData' (DV "pair" [PVar' "a",PVar' "b"])) + (NaryOp Sum [Var "a", Var "b"])] + + +match7 :: Text +match7 = unlines + ["match (-2.0,1.0). pair(real,prob):" + ," (a,b): normal(a,b)" + ] + +match7AST :: AST' Text +match7AST = Case (Ann + (Pair + (App (Var "negate") (ULiteral (Prob 2.0))) + (ULiteral (Prob 1.0))) + (TypeApp "pair" [TypeVar "real",TypeVar "prob"])) + [Branch' (PData' (DV "pair" [PVar' "a",PVar' "b"])) + (App (App (Var "normal") (Var "a")) (Var "b"))] + +testMatches :: Test +testMatches = test + [ testParse match1 match1AST + , testParse match2 match2AST + , testParse match3 match3AST + , testParse match4 match4AST + , testParse match5 match5AST + , testParse match6 match6AST + , testParse match7 match7AST + ] + +ann1 :: Text +ann1 = "5. nat" + +ann1AST :: AST' Text +ann1AST = Ann (ULiteral (Nat 5)) (TypeVar "nat") + +ann2 :: Text +ann2 = "(2,3). pair(a,b)" + +ann2AST :: AST' Text +ann2AST = + Ann (Pair (ULiteral (Nat 2)) (ULiteral (Nat 3))) + (TypeApp "pair" [TypeVar "a", TypeVar "b"]) + +ann3 :: Text +ann3 = "f. a -> measure(a)" + +ann3AST :: AST' Text +ann3AST = + Ann (Var "f") (TypeFun (TypeVar "a") + (TypeApp "measure" [(TypeVar "a")])) + +testAnn :: Test +testAnn = test + [ testParse ann1 ann1AST + , testParse ann2 ann2AST + , testParse ann3 ann3AST + ] + +expect1 :: Text +expect1 = unlines + ["expect x normal(0,1):" + ," 1" + ] + +expect1AST :: AST' Text +expect1AST = Expect "x" (App (App (Var "normal") + (ULiteral (Nat 0))) + (ULiteral (Nat 1))) + (ULiteral (Nat 1)) + +expect2 :: Text +expect2 = unlines + ["expect x normal(0,1):" + ," unsafeProb(x*x)" + ] + +expect2AST :: AST' Text +expect2AST = Expect "x" (App (App (Var "normal") + (ULiteral (Nat 0))) + (ULiteral (Nat 1))) + (App (Var "unsafeProb") + (NaryOp Prod [Var "x", Var "x"])) + +testExpect :: Test +testExpect = test + [ testParse expect1 expect1AST + , testParse expect2 expect2AST + ] + +array1 :: Text +array1 = unlines + ["array x of 12:" + ," x + 1" + ] + +array1AST :: AST' Text +array1AST = Array "x" (ULiteral (Nat 12)) + (NaryOp Sum [Var "x", ULiteral (Nat 1)]) + +array2 :: Text +array2 = "2 + x[3][4]" + +array2AST :: AST' Text +array2AST = NaryOp Sum + [ ULiteral (Nat 2) + , Index (Index (Var "x") + (ULiteral (Nat 3))) + (ULiteral (Nat 4)) + ] + +array3 :: Text +array3 = "[4, 0, 9]" + +array3AST :: AST' Text +array3AST = ArrayLiteral [ ULiteral (Nat 4) + , ULiteral (Nat 0) + , ULiteral (Nat 9) + ] + +testArray :: Test +testArray = test + [ testParse array1 array1AST + , testParse array2 array2AST + , testParse array3 array3AST + ] + +easyRoad1 :: Text +easyRoad1 = unlines + ["noiseT <~ uniform(3, 8)" + ,"noiseE <~ uniform(1, 4)" + ,"x1 <~ normal( 0, noiseT)" + ,"m1 <~ normal(x1, noiseE)" + ,"x2 <~ normal(x1, noiseT)" + ,"m2 <~ normal(x2, noiseE)" + ,"return ((m1, m2), (noiseT, noiseE))" + ] + +-- works in lax mode +easyRoad2 :: Text +easyRoad2 = unlines + ["noiseT_ <~ uniform(3, 8)" + ,"noiseE_ <~ uniform(1, 4)" + ,"noiseT = unsafeProb(noiseT_)" + ,"noiseE = unsafeProb(noiseE_)" + ,"x1 <~ normal(0, noiseT)" + ,"m1 <~ normal(x1, noiseE)" + ,"x2 <~ normal(x1, noiseT)" + ,"m2 <~ normal(x2, noiseE)" + ,"return ((m1, m2), (noiseT, noiseE)). measure(pair(pair(real,real),pair(prob,prob)))" + ] + + +easyRoadAST :: AST' Text +easyRoadAST = + Bind "noiseT" (App (App (Var "uniform") + (ULiteral (Nat 3))) + (ULiteral (Nat 8))) + (Bind "noiseE" (App (App (Var "uniform") + (ULiteral (Nat 1))) + (ULiteral (Nat 4))) + (Bind "x1" (App (App (Var "normal") + (ULiteral (Nat 0))) + (Var "noiseT")) + (Bind "m1" (App (App (Var "normal") + (Var "x1")) + (Var "noiseE")) + (Bind "x2" (App (App (Var "normal") + (Var "x1")) + (Var "noiseT")) + (Bind "m2" (App (App (Var "normal") + (Var "x2")) + (Var "noiseE")) + (Dirac + (Pair + (Pair + (Var "m1") + (Var "m2")) + (Pair + (Var "noiseT") + (Var "noiseE"))))))))) + +testRoadmap :: Test +testRoadmap = test + [ testParse easyRoad1 easyRoadAST + ] + + + +allTests :: Test +allTests = test + [ testIfs + , testLams + , testLets + , testBinds + , testMatches + , testAnn + , testExpect + , testArray + , testRoadmap + ] + + diff --git a/haskell/Tests/RoundTrip.hs b/haskell/Tests/RoundTrip.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/RoundTrip.hs @@ -0,0 +1,1418 @@ +{-# LANGUAGE NoImplicitPrelude + , DataKinds + , TypeOperators + , TypeFamilies + , ScopedTypeVariables + , FlexibleContexts + , MultiParamTypeClasses + , FunctionalDependencies + , TypeSynonymInstances + , GADTs + , FlexibleInstances + , FlexibleContexts + , ConstraintKinds + #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Tests.RoundTrip where + +import Prelude ((.), ($), asTypeOf, String, FilePath, Show(..), (++), Bool(..), concat) +import qualified Prelude +import qualified Data.List.NonEmpty as L +import Data.Ratio +import qualified Data.Text.Utf8 as IO + +import Language.Hakaru.Syntax.Prelude +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Pretty.Concrete (pretty) +import Language.Hakaru.Syntax.AST (Term, PrimOp(..)) +import Language.Hakaru.Syntax.AST.Transforms +import Language.Hakaru.Syntax.ABT (ABT, TrivialABT(..)) +import Language.Hakaru.Expect (total) +import Language.Hakaru.Inference (priorAsProposal, mcmc, mh) +import Language.Hakaru.Types.Sing +import System.IO +import System.Directory +import Control.Monad (mapM_, Monad(return)) +import Data.Foldable (null) +import Data.List (intercalate) + +import qualified Data.Text as Text +import Test.HUnit hiding ((~:), test) +import qualified Test.HUnit as HUnit +import Tests.TestTools hiding (testStriv, testSStriv, testConcreteFiles) +import qualified Tests.TestTools as Tools +import Tests.Models + (uniform_0_1, normal_0_1, gamma_1_1, + uniformC, normalC, beta_1_1, t4, t4', norm, unif2) + +unsafeSuperpose + :: (ABT Term abt) + => [(abt '[] 'HProb, abt '[] ('HMeasure a))] + -> abt '[] ('HMeasure a) +unsafeSuperpose = superpose . L.fromList + + +class IsTestGroup t where + test :: [t] -> t + +class IsTest' ta t | ta -> t, t -> ta where + (~:) :: String -> ta -> t + +class IsTestAssertion ta where + testStriv + :: TrivialABT Term '[] a + -> ta + + testSStriv + :: [(TrivialABT Term '[] a)] + -> TrivialABT Term '[] a + -> ta + + testConcreteFiles + :: FilePath + -> FilePath + -> ta + + +instance IsTestGroup Test where test = HUnit.test; +instance IsTest' Assertion Test where (~:) = (HUnit.~:) +instance IsTestAssertion Assertion where + testStriv = Tools.testStriv; testSStriv = Tools.testSStriv; testConcreteFiles = Tools.testConcreteFiles + + +data SaveInput = forall a . TestInput [TrivialABT Term '[] a] (TrivialABT Term '[] a) | DoNothing +newtype SaveTests = SaveTests { runSaveTests :: IO () } + +instance IsTestAssertion SaveInput where + testStriv = TestInput [] + testSStriv = TestInput + testConcreteFiles _ _ = DoNothing + +instance IsTestGroup SaveTests where + test = SaveTests . mapM_ runSaveTests + +instance IsTest' SaveInput SaveTests where + (~:) _ DoNothing = SaveTests (return ()) + (~:) tnm (TestInput xs r) = + let go (s,x) = do + createDirectoryIfMissing True $ intercalate "/" dn + (IO.writeFile fn . Text.pack . show . pretty . expandTransformations) x + where + dn = ["tests", "RoundTrip"] + fn = intercalate "/" $ dn ++ [concat [tnm, if null s then "" else ".", s, ".hk"]] + + xs' = case xs of + [] -> [("",r)] + _ -> ("expected",r) : Prelude.zip (Prelude.map Prelude.show [0..]) xs + in SaveTests $ mapM_ go xs' + +type IsTest ta t = (IsTest' ta t, IsTestGroup t, IsTestAssertion ta) + + +testMeasureUnit :: IsTest ta t => t +testMeasureUnit = test [ + "t1,t5" ~: testSStriv [t1,t5] (weight half), + "t10" ~: testSStriv [t10] (reject sing), + "t11,t22" ~: testSStriv [t11,t22] (dirac unit), + "t12" ~: testSStriv [] t12, + "t20" ~: testSStriv [t20] (lam $ \y -> weight (y * half)), + "t24" ~: testSStriv [t24] t24', + "t25" ~: testSStriv [t25] t25', + "t44Add" ~: testSStriv [t44Add] t44Add', + "t44Mul" ~: testSStriv [t44Mul] t44Mul', + "t53" ~: testSStriv [t53,t53'] t53'', + "t54" ~: testStriv t54, + "t55" ~: testSStriv [t55] t55', + "t56" ~: testSStriv [t56,t56'] t56'', + "t57" ~: testSStriv [t57] t57', + "t58" ~: testSStriv [t58] t58', + "t59" ~: testStriv t59, + "t60" ~: testSStriv [t60,t60'] t60'', + "t62" ~: testSStriv [t62] t62', + "t63" ~: testSStriv [t63] t63', + "t64" ~: testSStriv [t64,t64'] t64'', + "t65" ~: testSStriv [t65] t65', + "t77" ~: testSStriv [] t77 + ] + +testMeasureProb :: IsTest ta t => t +testMeasureProb = test [ + "t2" ~: testSStriv [t2] (unsafeProb <$> uniform zero one), + "t26" ~: testSStriv [t26] (dirac half), + "t30" ~: testSStriv [] t30, + "t33" ~: testSStriv [] t33, + "t34" ~: testSStriv [t34] (dirac (prob_ 3)), + "t35" ~: testSStriv [] t35, + "t35'" ~: testSStriv [] t35', + "t38" ~: testSStriv [] t38, + "t42" ~: testSStriv [t42] (dirac one), + "t49" ~: testSStriv [] t49, + "t61" ~: testSStriv [t61] t61', + "t66" ~: testSStriv [] t66, + "t67" ~: testSStriv [] t67, + "t69x" ~: testSStriv [t69x] (dirac $ prob_ 1.5), + "t69y" ~: testSStriv [t69y] (dirac $ prob_ 3.5) + ] + +testMeasureReal :: IsTest ta t => t +testMeasureReal = test + [ "t3" ~: testSStriv [] t3 + , "t6" ~: testSStriv [t6'] t6 + , "t7" ~: testSStriv [t7] t7' + , "t7n" ~: testSStriv [t7n] t7n' + , "t8'" ~: testSStriv [t8'] (lam $ \s1 -> + lam $ \s2 -> + normal zero (sqrt $ (s2 ^ (nat_ 2) + s1 ^ (nat_ 2)))) + , "t9" ~: testSStriv [t9] (unsafeSuperpose [(prob_ 2, uniform (real_ 3) (real_ 7))]) + , "t13" ~: testSStriv [t13] t13' + , "t14" ~: testSStriv [t14] t14' + , "t21" ~: testStriv t21 + , "t28" ~: testSStriv [] t28 + , "t31" ~: testSStriv [] t31 + , "t36" ~: testSStriv [] t36 + , "t37" ~: testSStriv [] t37 + , "t39" ~: testSStriv [] t39 + , "t40" ~: testSStriv [] t40 + , "t43" ~: testSStriv [t43, t43'] t43'' + , "t46" ~: testSStriv [] t46 + , "t45" ~: testSStriv [t47] t45 + , "t50" ~: testStriv t50 + , "t51" ~: testStriv t51 + , "t68" ~: testStriv t68 + , "t68'" ~: testStriv t68' + , "t70a" ~: testSStriv [t70a] (uniform one (real_ 3)) + , "t71a" ~: testSStriv [t71a] (uniform one (real_ 3)) + , "t72a" ~: testSStriv [t72a] (withWeight half $ uniform one (real_ 2)) + , "t73a" ~: testSStriv [t73a] (reject sing) + , "t74a" ~: testSStriv [t74a] (reject sing) + , "t70b" ~: testSStriv [t70b] (reject sing) + , "t71b" ~: testSStriv [t71b] (reject sing) + , "t72b" ~: testSStriv [t72b] (withWeight half $ uniform (real_ 2) (real_ 3)) + , "t73b" ~: testSStriv [t73b] (uniform one (real_ 3)) + , "t74b" ~: testSStriv [t74b] (uniform one (real_ 3)) + , "t70c" ~: testSStriv [t70c] (uniform one (real_ 3)) + , "t71c" ~: testSStriv [t71c] (uniform one (real_ 3)) + , "t72c" ~: testSStriv [t72c] (withWeight half $ uniform one (real_ 2)) + , "t73c" ~: testSStriv [t73c] (reject sing) + , "t74c" ~: testSStriv [t74c] (reject sing) + , "t70d" ~: testSStriv [t70d] (reject sing) + , "t71d" ~: testSStriv [t71d] (reject sing) + , "t72d" ~: testSStriv [t72d] (withWeight half $ uniform (real_ 2) (real_ 3)) + , "t73d" ~: testSStriv [t73d] (uniform one (real_ 3)) + , "t74d" ~: testSStriv [t74d] (uniform one (real_ 3)) + , "t76" ~: testStriv t76 + , "t78" ~: testSStriv [t78] t78' + , "t79" ~: testSStriv [t79] (dirac one) + , "t80" ~: testStriv t80 + , "t81" ~: testSStriv [] t81 + -- TODO, "kalman" ~: testStriv kalman + --, "seismic" ~: testSStriv [] seismic + , "lebesgue1" ~: testSStriv [] (lebesgue >>= \x -> if_ ((real_ 42) < x) (dirac x) (reject sing)) + , "lebesgue2" ~: testSStriv [] (lebesgue >>= \x -> if_ (x < (real_ 42)) (dirac x) (reject sing)) + , "lebesgue3" ~: testSStriv [lebesgue >>= \x -> if_ (x < (real_ 42) && (real_ 40) < x) (dirac x) (reject sing)] + (withWeight (prob_ $ 2) $ uniform (real_ 40) (real_ 42)) + , "testexponential" ~: testStriv testexponential + , "testcauchy" ~: testStriv testCauchy + , "exceptionLebesgue" ~: testSStriv [lebesgue >>= \x -> dirac (if_ (x == (real_ 3)) one x)] lebesgue + , "exceptionUniform" ~: testSStriv [uniform (real_ 2) (real_ 4) >>= \x -> + dirac (if_ (x == (real_ 3)) one x) + ] (uniform (real_ 2) (real_ 4)) + -- TODO "two_coins" ~: testStriv two_coins -- needs support for lists + ] + +testMeasureNat :: IsTest ta t => t +testMeasureNat = test + [ "size" ~: testConcreteFiles "tests/RoundTrip/size.0.hk" "tests/RoundTrip/size.expected.hk" + ] + +testMeasureInt :: IsTest ta t => t +testMeasureInt = test + [ "t75" ~: testStriv t75 + , "t75_hakaru" ~: testConcreteFiles "tests/t75_in.hk" "tests/t75_out.hk" + , "t75'" ~: testStriv t75' + , "t83" ~: testSStriv [t83] t83' + -- Jacques wrote: "bug: [simp_pw_equal] implicitly assumes the ambient measure is Lebesgue" + , "exceptionCounting" ~: testSStriv [] (counting >>= \x -> + if_ (x == (int_ 3)) + (dirac one) + (dirac x)) + , "exceptionSuperpose" ~: testSStriv + [(unsafeSuperpose [ (third, dirac (int_ 2)) + , (third, dirac (int_ 3)) + , (third, dirac (int_ 4)) + ] `asTypeOf` counting) >>= \x -> + dirac (if_ (x == (int_ 3)) one x)] + (unsafeSuperpose [ (third, dirac (int_ 2)) + , (third, dirac one) + , (third, dirac (int_ 4)) + ]) + ] + +testMeasurePair :: IsTest ta t => t +testMeasurePair = test [ + "t4" ~: testSStriv [t4] t4', + "t8" ~: testSStriv [] t8, + "t23" ~: testSStriv [t23] t23', + "t48" ~: testStriv t48, + "t52" ~: testSStriv [] t52, + "dup" ~: testSStriv [dup normal_0_1] (liftM2 pair + (normal zero one) + (normal zero one)), + "norm" ~: testSStriv [] norm, + "norm_nox" ~: testSStriv [norm_nox] (normal zero (sqrt (prob_ 2))), + "norm_noy" ~: testSStriv [norm_noy] (normal zero one), + "flipped_norm" ~: testSStriv [swap <$> norm] flipped_norm, + "priorProp" ~: testSStriv [lam (priorAsProposal norm)] + (lam $ \x -> unpair x $ \x0 x1 -> + unsafeSuperpose [(half, normal zero + (sqrt (prob_ 2)) >>= \y -> + dirac (pair x0 y)), + (half, normal_0_1 >>= \y -> + dirac (pair y x1))]), + "mhPriorProp" ~: testSStriv [testMHPriorProp] testPriorProp', + "unif2" ~: testStriv unif2, + "easyHMM" ~: testStriv easyHMM, + "testMCMCPriorProp" ~: testStriv testMCMCPriorProp + ] + +testOther :: IsTest ta t => t +testOther = test [ + "t82" ~: testSStriv [t82] t82', + "testRoadmapProg1" ~: testStriv rmProg1, + "testKernel" ~: testSStriv [testKernel] testKernel2 + --"testFalseDetection" ~: testStriv (lam seismicFalseDetection), + --"testTrueDetection" ~: testStriv (lam2 seismicTrueDetection) + --"testTrueDetectionL" ~: testStriv tdl, + --"testTrueDetectionR" ~: testStriv tdr + ] + +allTests :: IsTest ta t => t +allTests = test + [ testMeasureUnit + , testMeasureProb + , testMeasureReal + , testMeasurePair + , testMeasureNat + , testMeasureInt + , testOther + ] + +save_allTests :: IO () +save_allTests = runSaveTests allTests + +---------------------------------------------------------------- +-- In Maple, should 'evaluate' to "\c -> 1/2*c(Unit)" +t1 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t1 = uniform_0_1 >>= \x -> weight (unsafeProb x) + +t2 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t2 = beta_1_1 + +t3 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t3 = normal zero (prob_ 10) + +-- t5 is "the same" as t1. +t5 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t5 = weight half >> dirac unit + +t6, t6' :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t6 = dirac (real_ 5) +t6' = unsafeSuperpose [(one, dirac (real_ 5))] + +t7,t7', t7n,t7n' :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t7 = uniform_0_1 >>= \x -> weight (unsafeProb (x+one)) >> dirac (x*x) +t7' = uniform_0_1 >>= \x -> unsafeSuperpose [(unsafeProb (x+one), dirac (x^(nat_ 2)))] +t7n = + uniform (negate one) zero >>= \x -> + weight (unsafeProb (x+one)) >> + dirac (x*x) +t7n' = + uniform (real_ (-1)) zero >>= \x -> + unsafeSuperpose [(unsafeProb (x + one), dirac (x^(nat_ 2)))] + +-- For sampling efficiency (to keep importance weights at or close to 1), +-- t8 below should read back to uses of "normal", not uses of "lebesgue" +-- then "weight". +t8 :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HReal 'HReal)) +t8 = normal zero (prob_ 10) >>= \x -> normal x (prob_ 20) >>= \y -> dirac (pair x y) + +-- Normal is conjugate to normal +t8' :: (ABT Term abt) + => abt '[] ('HProb ':-> 'HProb ':-> 'HMeasure 'HReal) +t8' = + lam $ \s1 -> + lam $ \s2 -> + normal zero s1 >>= \x -> + normal x s2 + +t9 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t9 = + lebesgue >>= \x -> + weight (if_ ((real_ 3) < x && x < (real_ 7)) half zero) >> + dirac x + +t10 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t10 = weight zero + +t11 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t11 = weight one + +t12 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t12 = weight (prob_ 2) + +t13,t13' :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t13 = bern ((prob_ 3)/(prob_ 5)) >>= \b -> dirac (if_ b (real_ 37) (real_ 42)) +t13' = unsafeSuperpose + [ (prob_ $ 3 % 5, dirac (real_ 37)) + , (prob_ $ 2 % 5, dirac (real_ 42)) + ] + +t14,t14' :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t14 = + bern ((prob_ 3)/(prob_ 5)) >>= \b -> + if_ b t13 (bern ((prob_ 2)/(prob_ 7)) >>= \b' -> + if_ b' (uniform (real_ 10) (real_ 12)) (uniform (real_ 14) (real_ 16))) +t14' = unsafeSuperpose + [ (prob_ $ 9 % 25, dirac (real_ 37)) + , (prob_ $ 6 % 25, dirac (real_ 42)) + , (prob_ $ 4 % 35, uniform (real_ 10) (real_ 12)) + , (prob_ $ 2 % 7 , uniform (real_ 14) (real_ 16)) + ] + +t20 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure HUnit) +t20 = lam $ \y -> uniform_0_1 >>= \x -> weight (unsafeProb x * y) + +t21 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HReal) +t21 = mcmc (lam $ \x -> normal x one) (normal zero (prob_ 5)) + +t22 :: (ABT Term abt) => abt '[] ('HMeasure HUnit) +t22 = bern half >> dirac unit + +-- was called bayesNet in Nov.06 msg by Ken for exact inference +t23, t23' :: (ABT Term abt) => abt '[] ('HMeasure (HPair HBool HBool)) +t23 = + bern half >>= \a -> + bern (if_ a ((prob_ 9)/(prob_ 10)) ((prob_ 1)/(prob_ 10))) >>= \b -> + bern (if_ a ((prob_ 9)/(prob_ 10)) ((prob_ 1)/(prob_ 10))) >>= \c -> + dirac (pair b c) +t23' = unsafeSuperpose + [ ((prob_ $ 41 % 100), dirac (pair true true)) + , ((prob_ $ 9 % 100), dirac (pair true false)) + , ((prob_ $ 9 % 100), dirac (pair false true)) + , ((prob_ $ 41 % 100), dirac (pair false false)) + ] + +t24,t24' :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure HUnit) +t24 = + lam $ \x -> + uniform_0_1 >>= \y -> + uniform_0_1 >>= \z -> + weight (x * exp (cos y) * unsafeProb z) +t24' = + lam $ \x -> + withWeight (x * half) $ + uniform_0_1 >>= \y -> + weight (exp (cos y)) + +t25,t25' :: (ABT Term abt) + => abt '[] ('HProb ':-> 'HReal ':-> 'HMeasure HUnit) +t25 = + lam $ \x -> + lam $ \y -> + uniform_0_1 >>= \z -> + weight (x * exp (cos y) * unsafeProb z) +t25' = + lam $ \x -> + lam $ \y -> + weight (x * exp (cos y) * half) + +t26 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t26 = dirac (total t1) + +t28 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t28 = uniform zero one + +t30 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t30 = exp <$> uniform zero one + +t31 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t31 = uniform (real_ (-1)) one + +t33 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t33 = exp <$> t31 + +t34 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t34 = dirac (if_ ((real_ 2) < (real_ 4)) (prob_ 3) (prob_ 5)) + +t35, t35' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HProb) +t35 = lam $ \x -> dirac (if_ ((x `asTypeOf` log one) < (real_ 4)) (prob_ 3) (prob_ 5)) +t35' = lam $ \x -> if_ (x < (fromRational 4)) (dirac (fromRational 3)) (dirac (fromRational 5)) + +t36 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HProb) +t36 = lam (dirac . sqrt) + +t37 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HReal) +t37 = lam (dirac . recip) + +t38 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HProb) +t38 = lam (dirac . recip) + +t39 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HReal) +t39 = lam (dirac . log) + +t40 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HReal) +t40 = lam (dirac . log) + +t42 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t42 = dirac . total $ (unsafeProb <$> uniform zero (real_ 2)) + +t43, t43', t43'' :: (ABT Term abt) => abt '[] (HBool ':-> 'HMeasure 'HReal) +t43 = lam $ \b -> if_ b uniform_0_1 (fromProb <$> beta_1_1) +t43' = lam $ \b -> if_ b uniform_0_1 uniform_0_1 +t43'' = lam $ \_ -> uniform_0_1 + +t44Add, t44Add', t44Mul, t44Mul' + :: (ABT Term abt) => abt '[] ('HReal ':-> 'HReal ':-> 'HMeasure HUnit) +t44Add = lam $ \x -> lam $ \y -> weight (unsafeProb $ (x * x) + (y * y)) +t44Add' = lam $ \x -> lam $ \y -> weight (unsafeProb $ (x ^ (nat_ 2) + y ^ (nat_ 2))) +t44Mul = lam $ \x -> lam $ \y -> weight (unsafeProb $ (x * x * y * y)) +t44Mul' = lam $ \x -> lam $ \y -> weight (unsafeProb $ (x ^ (nat_ 2)) * (y ^ (nat_ 2))) + +-- t45, t46, t47 are all equivalent. +-- But t47 is worse than t45 and t46 because the importance weight generated by +-- t47 as a sampler varies between 0 and 1 whereas the importance weight generated +-- by t45 and t46 is always 1. In general it's good to reduce weight variance. +t45 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t45 = normal (real_ 4) (prob_ 5) >>= \x -> if_ (x < (real_ 3)) (dirac (x^(nat_ 2))) (dirac (x+(real_ (-1)))) + +t46 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t46 = normal (real_ 4) (prob_ 5) >>= \x -> dirac (if_ (x < (real_ 3)) (x*x) (x-one)) + +t47 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t47 = unsafeSuperpose + [ (one, normal (real_ 4) (prob_ 5) >>= \x -> if_ (x < (real_ 3)) (dirac (x*x)) (reject sing)) + , (one, normal (real_ 4) (prob_ 5) >>= \x -> if_ (x < (real_ 3)) (reject sing) (dirac (x-one))) + ] + +t48 :: (ABT Term abt) => abt '[] (HPair 'HReal 'HReal ':-> 'HMeasure 'HReal) +t48 = lam $ \x -> uniform (real_ (-5)) (real_ 7) >>= \w -> dirac ((fst x + snd x) * w) + +t49 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t49 = gamma (prob_ 0.01) (prob_ 0.35) + +t50 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t50 = uniform one (real_ 3) >>= \x -> normal one (unsafeProb x) + +t51 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t51 = t31 >>= \x -> normal x one + +-- Example 1 from Chang & Pollard's Conditioning as Disintegration +t52 :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HReal (HPair 'HReal 'HReal))) +t52 = + uniform_0_1 >>= \x -> + uniform_0_1 >>= \y -> + dirac (pair (max y x) (pair x y)) + +t53, t53', t53'' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t53 = + lam $ \x -> + unsafeSuperpose + [ (one, unsafeSuperpose + [ (one, + if_ (zero < x) + (if_ (x < one) (dirac unit) (reject sing)) + (reject sing)) + ]) + , (one, if_ false (dirac unit) (reject sing)) + ] +t53' = + lam $ \x -> + unsafeSuperpose + [ (one, + if_ (zero < x) + (if_ (x < one) (dirac unit) (reject sing)) + (reject sing)) + , (one, if_ false (dirac unit) (reject sing)) + ] +t53'' = + lam $ \x -> + if_ (zero < x && x < one) (dirac unit) (reject sing) + +t54 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t54 = + lam $ \x0 -> + ( dirac x0 >>= \x1 -> + (negate <$> uniform_0_1) >>= \x2 -> + dirac (x1 + x2) + ) >>= \x1 -> + ( ( (dirac zero >>= \x2 -> + dirac x1 >>= \x3 -> + dirac (x2 < x3) + ) >>= \x2 -> + if_ x2 + (recip <$> dirac x1) + (dirac zero) + ) >>= \x2 -> + weight (unsafeProb x2) + ) >> + (log <$> dirac (unsafeProb x1)) >>= \x3 -> + (negate <$> dirac x3) >>= \x4 -> + ( + (dirac zero >>= \x5 -> + dirac x4 >>= \x6 -> + dirac (x5 < x6) + ) >>= \x5 -> + if_ x5 + ( (dirac x4 >>= \x6 -> + dirac one >>= \x7 -> + dirac (x6 < x7) + ) >>= \x6 -> + if_ x6 (dirac one) (dirac zero) + ) + (dirac zero) + ) >>= \x5 -> + weight (unsafeProb x5) + +t55, t55' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t55 = + lam $ \t -> + uniform_0_1 >>= \x -> + if_ (x < t) (dirac unit) (reject sing) +t55' = + lam $ \t -> + if_ (t < zero) (reject sing) $ + if_ (t < one) (weight (unsafeProb t)) $ + dirac unit + +t56, t56', t56'' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t56 = + lam $ \x0 -> + ( dirac x0 >>= \x1 -> + (negate <$> uniform_0_1) >>= \x2 -> + dirac (x1 + x2) + ) >>= \x1 -> + ( (dirac zero >>= \x2 -> + dirac x1 >>= \x3 -> + dirac (x2 < x3) + ) >>= \x2 -> + if_ x2 + ( (dirac x1 >>= \x3 -> + dirac one >>= \x4 -> + dirac (x3 < x4) + ) >>= \x3 -> + if_ x3 (dirac one) (dirac zero)) + (dirac zero) + ) >>= \x2 -> + withWeight (unsafeProb x2) (dirac unit) +t56' = + lam $ \x0 -> + uniform_0_1 >>= \x1 -> + if_ (x0 - one < x1 && x1 < x0) + (dirac unit) + (reject sing) +t56'' = + lam $ \t -> + if_ (t <= zero) (reject sing) $ + if_ (t <= one) (weight (unsafeProb t)) $ + if_ (t <= (real_ 2)) (weight (unsafeProb ((real_ 2) + t * negate one))) $ + reject sing + +t57, t57' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t57 = lam $ \t -> unsafeSuperpose + [ (one, if_ (t < one) (dirac unit) (reject sing)) + , (one, if_ (zero < t) (dirac unit) (reject sing)) ] +t57' = lam $ \t -> + if_ (t < one && zero < t) (weight (prob_ 2)) (dirac unit) + +t58, t58' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t58 = lam $ \t -> unsafeSuperpose + [ (one, if_ (zero < t && t < (real_ 2)) (dirac unit) (reject sing)) + , (one, if_ (one < t && t < (real_ 3)) (dirac unit) (reject sing)) ] +t58' = lam $ \t -> + if_ (if_ (zero < t) (t < (real_ 2)) false) + (if_ (if_ (one < t) (t < (real_ 3)) false) + (weight (prob_ 2)) + (dirac unit)) + (if_ (if_ (one < t) (t < (real_ 3)) false) + (dirac unit) + (reject sing)) + +t59 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t59 = + lam $ \x0 -> + ((recip <$> uniform_0_1) >>= \x1 -> + (((dirac zero >>= \x2 -> + dirac x1 >>= \x3 -> + dirac (x2 < x3)) >>= \x2 -> + if_ x2 + (dirac x1) + (negate <$> dirac x1)) >>= \x2 -> + weight (unsafeProb x2) ) >> + dirac x0 >>= \x3 -> + dirac x1 >>= \x4 -> + dirac (x3 * x4)) >>= \x1 -> + (dirac x1 >>= \x2 -> + (negate <$> uniform_0_1) >>= \x3 -> + dirac (x2 + x3)) >>= \x2 -> + ((dirac zero >>= \x3 -> + dirac x2 >>= \x4 -> + dirac (x3 < x4)) >>= \x3 -> + if_ x3 + ((dirac x2 >>= \x4 -> + dirac one >>= \x5 -> + dirac (x4 < x5)) >>= \x4 -> + if_ x4 (dirac one) (dirac zero)) + (dirac zero)) >>= \x3 -> + weight (unsafeProb x3) + +t60,t60',t60'' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t60 = + lam $ \x0 -> + (((uniform_0_1 >>= \x1 -> + uniform_0_1 >>= \x2 -> + dirac (x1 + x2)) >>= \x1 -> + dirac (recip x1)) >>= \x1 -> + (((dirac zero >>= \x2 -> + dirac x1 >>= \x3 -> + dirac (x2 < x3)) >>= \x2 -> + if_ x2 + (dirac x1) + (negate <$> dirac x1)) >>= \x2 -> + weight (unsafeProb x2) ) >> + dirac x0 >>= \x3 -> + dirac x1 >>= \x4 -> + dirac (x3 * x4)) >>= \x1 -> + ((dirac zero >>= \x2 -> + dirac x1 >>= \x3 -> + dirac (x2 < x3)) >>= \x2 -> + if_ x2 + ((dirac x1 >>= \x3 -> + dirac one >>= \x4 -> + dirac (x3 < x4)) >>= \x3 -> + if_ x3 (dirac one) (dirac zero)) + (dirac zero)) >>= \x2 -> + weight (unsafeProb x2) +t60' = + lam $ \x0 -> + uniform_0_1 >>= \x1 -> + uniform_0_1 >>= \x2 -> + if_ (if_ (zero < x0 / (x2 + x1)) + (x0 / (x2 + x1) < one) + false) + (weight ((unsafeProb (x2 + x1)) ^^ negate one) ) + (reject sing) +t60'' = + lam $ \x0 -> + uniform_0_1 >>= \x1 -> + uniform_0_1 >>= \x2 -> + if_ (if_ (zero < x0 / (x2 + x1)) + (x0 / (x2 + x1) < one) + false) + (weight (recip (unsafeProb (x2 + x1))) ) + (reject sing) + +t61, t61' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HProb) +t61 = lam $ \x -> if_ (x < zero) (dirac zero) $ dirac $ unsafeProb $ recip x +t61'= lam $ \x -> if_ (x < zero) (dirac zero) $ dirac $ unsafeProb $ recip x + +---- "Special case" of t56 +t62, t62' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HReal ':-> 'HMeasure HUnit) +t62 = lam $ \t -> + lam $ \x -> + uniform_0_1 >>= \y -> + if_ (zero < t/x - y && t/x - y < one) + (dirac unit) + (reject sing) +t62'= lam $ \t -> + lam $ \x -> + if_ (t/x <= zero) (reject sing) $ + if_ (t/x <= one) (weight (unsafeProb (t/x))) $ + if_ (t/x <= (real_ 2)) (weight (unsafeProb ((real_ 2)-t/x))) $ + reject sing + +---- "Scalar multiple" of t62 +t63, t63' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t63 = lam $ \t -> + uniform_0_1 >>= \x -> + uniform_0_1 >>= \y -> + if_ (zero < t/x - y && t/x - y < one) + (weight (recip (unsafeProb x))) + (reject sing) +t63'= lam $ \t -> + uniform_0_1 >>= \x -> + if_ (t/x <= zero) (reject sing) $ + if_ (t/x <= one) (weight (unsafeProb (t/x) / unsafeProb x)) $ + if_ (t/x <= (real_ 2)) (weight (unsafeProb ((real_ 2)-t/x) / unsafeProb x)) $ + reject sing + +-- Density calculation for (Exp (Log StdRandom)) and StdRandom +t64, t64', t64'' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t64 = lam $ \x0 -> + (((dirac zero >>= \x1 -> + dirac x0 >>= \x2 -> + dirac (x1 < x2)) >>= \x1 -> + if_ x1 + (recip <$> dirac x0) + (dirac zero)) >>= \x1 -> + weight (unsafeProb x1)) >> + (log <$> dirac (unsafeProb x0)) >>= \x2 -> + ((exp <$> dirac x2) >>= \x3 -> + weight x3) >> + (exp <$> dirac x2) >>= \x4 -> + ((dirac zero >>= \x5 -> + dirac x4 >>= \x6 -> + dirac (x5 < x6)) >>= \x5 -> + if_ x5 + ((dirac x4 >>= \x6 -> + dirac one >>= \x7 -> + dirac (x6 < x7)) >>= \x6 -> + if_ x6 (dirac one) (dirac zero)) + (dirac zero)) >>= \x5 -> + weight (unsafeProb x5) +t64' =lam $ \x0 -> + ((dirac zero >>= \x1 -> + dirac x0 >>= \x2 -> + dirac (x1 < x2)) >>= \x1 -> + if_ x1 + ((dirac x0 >>= \x2 -> + dirac one >>= \x3 -> + dirac (x2 < x3)) >>= \x2 -> + if_ x2 (dirac one) (dirac zero)) + (dirac zero)) >>= \x1 -> + weight (unsafeProb x1) +t64''=lam $ \x0 -> + if_ (zero < x0 && x0 < one) + (dirac unit) + (reject sing) + +-- Density calculation for (Add StdRandom (Exp (Neg StdRandom))). +-- Maple can integrate this but we don't simplify it for some reason. +t65, t65' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t65 = + lam $ \t -> + uniform_0_1 >>= \x -> + if_ (zero < t-x) + (let_ (unsafeProb (t-x)) $ \t_x -> + withWeight (recip t_x) $ + (if_ (zero < negate (log t_x) && negate (log t_x) < one) + (dirac unit) + (reject sing))) + (reject sing) +t65' = + lam $ \t -> + uniform_0_1 >>= \x-> + withWeight (if_ (real_ 0 < (log (unsafeProb (t + x * real_ (-1))) * real_ (-1)) && + x < (t * fromProb (exp (real_ 1)) + real_ (-1)) * fromProb (exp (real_ (-1))) && + x < t) + (unsafeProb (recip (x * real_ (-1) + t))) + (prob_ 0)) $ (dirac unit) + +half' :: (ABT Term abt) => abt '[] 'HReal +half' = half + +t66 :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t66 = dirac (sqrt $ prob_ 3 + (sqrt $ prob_ 3)) + +t67 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HReal ':-> 'HMeasure 'HProb) +t67 = lam $ \p -> lam $ \r -> dirac (exp (r * fromProb p)) + +t68 :: (ABT Term abt) + => abt '[] ('HProb ':-> 'HProb ':-> 'HReal ':-> 'HMeasure 'HReal) +t68 = + lam $ \x4 -> + lam $ \x5 -> + lam $ \x1 -> + lebesgue >>= \x2 -> + lebesgue >>= \x3 -> + withWeight (exp (negate (x2 - x3) * (x2 - x3) + * recip (fromProb ((fromRational 2) * exp (log x4 * (fromRational 2))))) + * recip x4 + * recip (exp (log ((fromRational 2) * pi) * half))) + (withWeight (exp (negate (x1 - x3) * (x1 - x3) + * recip (fromProb ((fromRational 2) * exp (log x5 * (fromRational 2))))) + * recip x5 + * recip (exp (log ((fromRational 2) * pi) * half))) + (withWeight (exp (negate x3 * x3 + * recip (fromProb ((fromRational 2) * exp (log x4 * (fromRational 2))))) + * recip x4 + * recip (exp (log ((fromRational 2) * pi) * half))) + (dirac x2))) + +t68' :: (ABT Term abt) => abt '[] ('HProb ':-> 'HReal ':-> 'HMeasure 'HReal) +t68' = lam $ \noise -> app (app t68 noise) noise + +t69x, t69y :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +t69x = dirac (integrate one (real_ 2) $ \x -> integrate (real_ 3) (real_ 4) $ \_ -> unsafeProb x) +t69y = dirac (integrate one (real_ 2) $ \_ -> integrate (real_ 3) (real_ 4) $ \y -> unsafeProb y) + +t70a, t71a, t72a, t73a, t74a :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t70a = uniform one (real_ 3) >>= \x -> if_ ((real_ 4) < x) (reject sing) (dirac x) +t71a = uniform one (real_ 3) >>= \x -> if_ ((real_ 3) < x) (reject sing) (dirac x) +t72a = uniform one (real_ 3) >>= \x -> if_ ((real_ 2) < x) (reject sing) (dirac x) +t73a = uniform one (real_ 3) >>= \x -> if_ (one < x) (reject sing) (dirac x) +t74a = uniform one (real_ 3) >>= \x -> if_ (zero < x) (reject sing) (dirac x) + +t70b, t71b, t72b, t73b, t74b :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t70b = uniform one (real_ 3) >>= \x -> if_ ((real_ 4) < x) (dirac x) (reject sing) +t71b = uniform one (real_ 3) >>= \x -> if_ ((real_ 3) < x) (dirac x) (reject sing) +t72b = uniform one (real_ 3) >>= \x -> if_ ((real_ 2) < x) (dirac x) (reject sing) +t73b = uniform one (real_ 3) >>= \x -> if_ (one < x) (dirac x) (reject sing) +t74b = uniform one (real_ 3) >>= \x -> if_ (zero < x) (dirac x) (reject sing) + +t70c, t71c, t72c, t73c, t74c :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t70c = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 4)) (dirac x) (reject sing) +t71c = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 3)) (dirac x) (reject sing) +t72c = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 2)) (dirac x) (reject sing) +t73c = uniform one (real_ 3) >>= \x -> if_ (x < one) (dirac x) (reject sing) +t74c = uniform one (real_ 3) >>= \x -> if_ (x < zero) (dirac x) (reject sing) + +t70d, t71d, t72d, t73d, t74d :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t70d = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 4)) (reject sing) (dirac x) +t71d = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 3)) (reject sing) (dirac x) +t72d = uniform one (real_ 3) >>= \x -> if_ (x < (real_ 2)) (reject sing) (dirac x) +t73d = uniform one (real_ 3) >>= \x -> if_ (x < one) (reject sing) (dirac x) +t74d = uniform one (real_ 3) >>= \x -> if_ (x < zero) (reject sing) (dirac x) + +t75 :: (ABT Term abt) => abt '[] ('HMeasure 'HNat) +t75 = gamma (prob_ 6) one >>= poisson + +t75' :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HNat) +t75' = lam $ \x -> gamma x one >>= poisson + +t76 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HReal) +t76 = + lam $ \x -> + lebesgue >>= \y -> + withWeight (unsafeProb (abs y)) $ + if_ (y < one) + (if_ (zero < y) + (if_ (x * y < one) + (if_ (zero < x * y) + (dirac (x * y)) + (reject sing)) + (reject sing)) + (reject sing)) + (reject sing) + +-- the (x * (-1)) below is an unfortunate artifact not worth fixing +t77 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure HUnit) +t77 = + lam $ \x -> + if_ (x < zero) + (weight (recip (exp x))) + (weight (exp x)) + +t78, t78' :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t78 = uniform zero (real_ 2) >>= \x2 -> withWeight (unsafeProb x2) (dirac x2) +t78' = beta (prob_ 2) one >>= \x -> dirac ((fromProb x) * (real_ 2)) + +-- what does this simplify to? +t79 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t79 = dirac (real_ 3) >>= \x -> dirac (if_ (x == (real_ 3)) one x) + +t80 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t80 = gamma_1_1 >>= \t -> normal zero t + +t81 :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +t81 = uniform zero pi + +t82 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HProb) +t82 = lam (densityUniform zero one) + +t82' :: (ABT Term abt) => abt '[] ('HReal ':-> 'HProb) +t82' = lam $ \x -> one + +t83 :: (ABT Term abt) => abt '[] ('HNat ':-> 'HMeasure 'HNat) +t83 = lam $ \k -> + plate k (\_ -> dirac (nat_ 1)) >>= \x -> + dirac (size x) + +t83' :: (ABT Term abt) => abt '[] ('HNat ':-> 'HMeasure 'HNat) +t83' = lam dirac + +-- Testing round-tripping of some other distributions +testexponential :: (ABT Term abt) => abt '[] ('HMeasure 'HProb) +testexponential = exponential third + +testCauchy :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +testCauchy = cauchy (real_ 5) (prob_ 3) + +testMCMCPriorProp + :: (ABT Term abt) + => abt '[] (HPair 'HReal 'HReal ':-> 'HMeasure (HPair 'HReal 'HReal)) +testMCMCPriorProp = mcmc (lam $ priorAsProposal norm) norm + +testMHPriorProp + :: (ABT Term abt) + => abt '[] + (HPair 'HReal 'HReal + ':-> 'HMeasure (HPair (HPair 'HReal 'HReal) 'HProb)) +testMHPriorProp = mh (lam $ priorAsProposal norm) norm + +testPriorProp' + :: (ABT Term abt) + => abt '[] + (HPair 'HReal 'HReal + ':-> 'HMeasure (HPair (HPair 'HReal 'HReal) 'HProb)) +testPriorProp' = + lam $ \old -> + unsafeSuperpose + [(half, + normal_0_1 >>= \x1 -> + dirac (pair (pair x1 (snd old)) + (exp + ( (x1 * negate one + (old `unpair` \x2 x3 -> x2)) + * ( (old `unpair` \x2 x3 -> x2) + + (old `unpair` \x2 x3 -> x3) * (negate (real_ 2)) + + x1) + * half)))) + , (half, + normal zero (sqrt (prob_ 2)) >>= \x1 -> + dirac (pair (pair (fst old) x1) + (exp + ( (x1 + (old `unpair` \x2 x3 -> x3) * negate one) + * ( (old `unpair` \x2 x3 -> x3) + + (old `unpair` \x2 x3 -> x2) * (negate (real_ 4)) + + x1) + * (negate (real_ 1))/(real_ 4))))) + ] + +dup :: (ABT Term abt, SingI a) + => abt '[] ('HMeasure a) + -> abt '[] ('HMeasure (HPair a a)) +dup m = let_ m (\m' -> liftM2 pair m' m') + +norm_nox :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +norm_nox = + normal_0_1 >>= \x -> + normal x one >>= \y -> + dirac y + +norm_noy :: (ABT Term abt) => abt '[] ('HMeasure 'HReal) +norm_noy = + normal_0_1 >>= \x -> + normal x one >> + dirac x + +flipped_norm :: (ABT Term abt) => abt '[] ('HMeasure (HPair 'HReal 'HReal)) +flipped_norm = + normal zero one >>= \x -> + normal x one >>= \y -> + dirac (pair y x) + +-- pull out some of the intermediate expressions for independent study +expr1 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HProb) +expr1 = + lam $ \x0 -> + (lam $ \_ -> + lam $ \x2 -> + lam $ \x3 -> + (lam $ \x4 -> + zero + + one + * (lam $ \x5 -> + (lam $ \x6 -> + zero + + exp (negate (x2 - zero) * (x2 - zero) / fromProb ((fromRational 2) * exp (log (fromRational 5) * (fromRational 2)))) + / (fromRational 5) + / exp (log ((fromRational 2) * pi) * half) + * (lam $ \x7 -> x7 `app` unit) `app` x6) + `app` (lam $ \_ -> + (lam $ \x7 -> + (lam $ \x8 -> x8 `app` x2) + `app` (lam $ \_ -> + (lam $ \x9 -> + (lam $ \x10 -> x10 `app` unit) + `app` (lam $ \x10 -> + (lam $ \x11 -> + (lam $ \x12 -> x12 `app` x2) + `app` (lam $ \_ -> + (lam $ \x13 -> x13 `app` pair x2 x10) `app` x11)) + `app` x9)) + `app` x7)) + `app` x5)) + `app` x4) + `app` (lam $ \x4 -> + (lam $ \x5 -> x5 `app` (x4 `unpair` \_ x7 -> x7)) `app` x3) + ) + `app` unit + `app` x0 + `app` (lam $ \_ -> one) + +expr2 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HReal ':-> 'HProb) +expr2 = + lam $ \x1 -> + lam $ \x2 -> + (lam $ \x3 -> + lam $ \x4 -> + lam $ \x5 -> + (lam $ \x6 -> + zero + + one + * (lam $ \x7 -> + (lam $ \x8 -> + zero + + exp (((negate x4) - x3) * (x4 - x3) / fromProb ((fromRational 2) * exp (log one * (fromRational 2)))) + / one + / exp (log ((fromRational 2) * pi) * half) + * (lam $ \x9 -> x9 `app` unit) `app` x8) + `app` (lam $ \_ -> + (lam $ \x9 -> + (lam $ \x10 -> x10 `app` x4) + `app` (lam $ \_ -> + (lam $ \x11 -> + (lam $ \x12 -> x12 `app` unit) + `app` (lam $ \x12 -> + (lam $ \x13 -> + (lam $ \x14 -> x14 `app` x4) + `app` (lam $ \_ -> + (lam $ \x15 -> x15 `app` pair x4 x12) `app` x13)) + `app` x11)) + `app` x9)) + `app` x7)) + `app` x6) + `app` (lam $ \x6 -> + (lam $ \x7 -> x7 `app` (x6 `unpair` \_ x9 -> x9)) `app` x5) + ) + `app` x1 + `app` x2 + `app` (lam $ \_ -> one) + +-- the one we need in testKernel +expr3 :: (ABT Term abt) + => abt '[] (d ':-> 'HProb) + -> abt '[] (d ':-> d ':-> 'HProb) + -> abt '[] d -> abt '[] d -> abt '[] 'HProb +expr3 x0 x1 x2 x3 = + let q = x0 `app` x3 + / x1 `app` x2 `app` x3 + * x1 `app` x3 `app` x2 + / x0 `app` x2 + in if_ (one < q) one q + +-- testKernel :: Sample IO ('HReal ':-> 'HMeasure 'HReal) +testKernel :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HReal) +testKernel = +-- Below is the output of testMcmc as of 2014-11-05 + let_ expr1 $ \x0 -> + let_ expr2 $ \x1 -> + lam $ \x2 -> + normal x2 one >>= \x3 -> + let_ (expr3 x0 x1 x2 x3) $ \x4 -> + bern x4 >>= \x5 -> + dirac (if_ x5 x3 x2) + +-- this should be equivalent to the above +testKernel2 :: (ABT Term abt) => abt '[] ('HReal ':-> 'HMeasure 'HReal) +testKernel2 = + lam $ \x2 -> + normal x2 one >>= \x3 -> + let q = exp(negate (real_ 1)/(real_ 50)*(x3-x2)*(x3+x2)) in + let_ (if_ (one < q) one q) $ \x4 -> + bern x4 >>= \x5 -> + dirac $ if_ x5 x3 x2 + +-- this comes from {Tests.Lazy,Examples.EasierRoadmap}.easierRoadmapProg1. It is the +-- program post-disintegrate, as passed to Maple to simplify +rmProg1 :: (ABT Term abt) => abt '[] + (HUnit + ':-> HPair 'HReal 'HReal + ':-> 'HMeasure (HPair 'HProb 'HProb)) +rmProg1 = + lam $ \_ -> + lam $ \x1 -> + x1 `unpair` \x2 x3 -> + withWeight one $ + withWeight one $ + unsafeSuperpose + [(one, + withWeight one $ + lebesgue >>= \x4 -> + unsafeSuperpose + [(one, + withWeight one $ + lebesgue >>= \x5 -> + withWeight one $ + lebesgue >>= \x6 -> + withWeight + ( exp (negate (x3 - x6) * (x3 - x6) + / (fromProb ((fromRational 2) * exp (log (unsafeProb x5) * (fromRational 2))))) + / unsafeProb x5 + / (exp (log ((fromRational 2) * pi) * half))) $ + withWeight one $ + lebesgue >>= \x7 -> + withWeight + ( exp (negate (x6 - x7) * (x6 - x7) + / (fromProb ((fromRational 2) * exp (log (unsafeProb x4) * (fromRational 2))))) + / (unsafeProb x4) + / (exp (log ((fromRational 2) * pi) * half))) $ + withWeight + ( exp (negate (x2 - x7) * (x2 - x7) + / (fromProb ((fromRational 2) * exp (log (unsafeProb x5) * (fromRational 2))))) + / unsafeProb x5 + / (exp (log ((fromRational 2) * pi) * half))) $ + withWeight + ( exp (negate x7 * x7 + / (fromProb ((fromRational 2) * exp (log (unsafeProb x4) * (fromRational 2))))) + / unsafeProb x4 + / (exp (log ((fromRational 2) * pi) * half))) $ + withWeight (recip (fromRational 3)) $ + unsafeSuperpose + [(one, + if_ (x5 < (real_ 4)) + (if_ (one < x5) + (withWeight (recip (prob_ 5)) $ + unsafeSuperpose + [(one, + if_ (x4 < (real_ 8)) + (if_ ((real_ 3) < x4) + (dirac (pair (unsafeProb x4) + (unsafeProb x5))) + (reject sing)) + (reject sing)) + , (one, reject sing)]) + (reject sing)) + (reject sing)) + , (one, reject sing)]) + , (one, reject sing)]) + , (one, reject sing)] + +-- this comes from Examples.EasierRoadmap.easierRoadmapProg4'. +-- TODO: this needs to be regenerated from original program +rmProg4 + :: (ABT Term abt) + => abt '[] + (HPair 'HReal 'HReal + ':-> HPair 'HProb 'HProb + ':-> 'HMeasure (HPair (HPair 'HProb 'HProb) 'HProb)) +rmProg4 = + lam $ \x0 -> + let_ (lam $ \x1 -> + (lam $ \x2 -> + lam $ \x3 -> + x3 `unpair` \x4 x5 -> + let_ one $ \x6 -> + let_ (let_ one $ \x7 -> + let_ (let_ one $ \x8 -> + let_ (let_ one $ \x9 -> + let_ (let_ one $ \x10 -> + let_ (let_ one $ \x11 -> + let_ (x2 `unpair` \x12 _ -> + x2 `unpair` \x14 _ -> + x2 `unpair` \x16 _ -> + x2 `unpair` \_ x19 -> + x2 `unpair` \_ x21 -> + x2 `unpair` \_ x23 -> + x2 `unpair` \x24 _ -> + x2 `unpair` \x26 _ -> + x2 `unpair` \_ x29 -> + x2 `unpair` \_ x31 -> + let_ (recip pi + * exp ((x12 * x14 * (fromProb x4 * fromProb x4) + * (fromRational 2) + + fromProb x4 * fromProb x4 * x16 * x19 + * (negate (fromRational 2)) + + x21 * x23 * (fromProb x4 * fromProb x4) + + fromProb x5 * fromProb x5 * (x24 * x26) + + fromProb x5 * fromProb x5 * (x29 * x31)) + * recip (fromProb x4 * fromProb x4 + * (fromProb x4 * fromProb x4) + + fromProb x5 * fromProb x5 + * (fromProb x4 * fromProb x4) + * (fromRational 3) + + fromProb x5 * fromProb x5 + * (fromProb x5 * fromProb x5)) + * (negate half)) + * exp (log (exp (log x4 * (fromRational 4)) + + exp (log x5 * (fromRational 2)) + * exp (log x4 * (fromRational 2)) + * (fromRational 3) + + exp (log x5 * (fromRational 4))) + * (negate half)) + * (fromRational 1)/(fromRational 10)) $ \x32 -> + let_ (let_ (recip (fromRational 3)) $ \x33 -> + let_ (let_ one $ \x34 -> + let_ (if_ (fromProb x5 < (fromRational 4)) + (if_ (one < fromProb x5) + (let_ (recip (fromRational 5)) $ \x35 -> + let_ (let_ one $ \x36 -> + let_ (if_ (fromProb x4 + < (fromRational 8)) + (if_ ((fromRational 3) + < fromProb x4) + (let_ (fromRational 5) $ \x37 -> + let_ (let_ (pair x4 x5) $ \x38 -> + pair (dirac x38) + (lam $ \x39 -> + x39 + `app` x38)) $ \x38 -> + pair (withWeight x37 $ + x38 `unpair` \x39 _ -> + x39) + (lam $ \x39 -> + zero + + x37 + * (x38 `unpair` \_ x41 -> + x41) + `app` x39)) + (pair (reject sing) + (lam $ \x37 -> + zero))) + (pair (reject sing) + (lam $ \x37 -> + zero))) $ \x37 -> + let_ one $ \x38 -> + let_ (pair (reject sing) + (lam $ \x39 -> + zero)) $ \x39 -> + pair (unsafeSuperpose [(x36, + x37 `unpair` \x40 x41 -> + x40), + (x38, + x39 `unpair` \x40 x41 -> + x40)]) + (lam $ \x40 -> + zero + + x36 + * (x37 `unpair` \x41 x42 -> + x42) + `app` x40 + + x38 + * (x39 `unpair` \x41 x42 -> + x42) + `app` x40)) $ \x36 -> + pair (withWeight x35 $ + x36 `unpair` \x37 x38 -> + x37) + (lam $ \x37 -> + zero + + x35 + * (x36 `unpair` \x38 x39 -> + x39) + `app` x37)) + (pair (reject sing) + (lam $ \x35 -> zero))) + (pair (reject sing) + (lam $ \x35 -> zero))) $ \x35 -> + let_ one $ \x36 -> + let_ (pair (reject sing) + (lam $ \x37 -> zero)) $ \x37 -> + pair (unsafeSuperpose [(x34, + x35 `unpair` \x38 x39 -> + x38), + (x36, + x37 `unpair` \x38 x39 -> + x38)]) + (lam $ \x38 -> + zero + + x34 + * (x35 `unpair` \x39 x40 -> x40) + `app` x38 + + x36 + * (x37 `unpair` \x39 x40 -> x40) + `app` x38)) $ \x34 -> + pair (withWeight x33 $ x34 `unpair` \x35 x36 -> x35) + (lam $ \x35 -> + zero + + x33 + * (x34 `unpair` \x36 x37 -> x37) + `app` x35)) $ \x33 -> + pair (withWeight x32 $ x33 `unpair` \x34 x35 -> x34) + (lam $ \x34 -> + zero + + x32 + * (x33 `unpair` \x35 x36 -> x36) + `app` x34)) $ \x12 -> + pair (withWeight x11 $ x12 `unpair` \x13 x14 -> x13) + (lam $ \x13 -> + zero + + x11 + * (x12 `unpair` \x14 x15 -> x15) `app` x13)) $ \x11 -> + let_ one $ \x12 -> + let_ (pair (reject sing) (lam $ \x13 -> zero)) $ \x13 -> + pair (unsafeSuperpose [(x10, x11 `unpair` \x14 x15 -> x14), + (x12, x13 `unpair` \x14 x15 -> x14)]) + (lam $ \x14 -> + zero + x10 * (x11 `unpair` \x15 x16 -> x16) `app` x14 + + x12 * (x13 `unpair` \x15 x16 -> x16) `app` x14)) $ \x10 -> + pair (withWeight x9 $ x10 `unpair` \x11 x12 -> x11) + (lam $ \x11 -> + zero + x9 * (x10 `unpair` \x12 x13 -> x13) `app` x11)) $ \x9 -> + let_ one $ \x10 -> + let_ (pair (reject sing) (lam $ \x11 -> zero)) $ \x11 -> + pair (unsafeSuperpose [(x8, x9 `unpair` \x12 x13 -> x12), + (x10, x11 `unpair` \x12 x13 -> x12)]) + (lam $ \x12 -> + zero + x8 * (x9 `unpair` \x13 x14 -> x14) `app` x12 + + x10 * (x11 `unpair` \x13 x14 -> x14) `app` x12)) $ \x8 -> + pair (withWeight x7 $ x8 `unpair` \x9 x10 -> x9) + (lam $ \x9 -> + zero + x7 * (x8 `unpair` \x10 x11 -> x11) `app` x9)) $ \x7 -> + pair (withWeight x6 $ x7 `unpair` \x8 x9 -> x8) + (lam $ \x8 -> zero + x6 * (x7 `unpair` \x9 x10 -> x10) `app` x8)) + `app` x0 + `app` x1 `unpair` \x2 x3 -> + x3 `app` (lam $ \x4 -> one)) $ \x1 -> + lam $ \x2 -> + (x2 `unpair` \x3 x4 -> + unsafeSuperpose [(half, + uniform (real_ 3) (real_ 8) >>= \x5 -> dirac (pair (unsafeProb x5) x4)), + (half, + uniform one (real_ 4) >>= \x5 -> + dirac (pair x3 (unsafeProb x5)))]) >>= \x3 -> + dirac (pair x3 (x1 `app` x3 / x1 `app` x2)) + + +pairReject + :: (ABT Term abt) + => abt '[] (HPair ('HMeasure 'HReal) 'HReal) +pairReject = + pair (reject (SMeasure SReal) >>= \_ -> dirac one) + (real_ 2) + +-- from a web question +-- these are mathematically equivalent, albeit at different types +chal1 :: (ABT Term abt) => abt '[] + ('HProb ':-> 'HReal ':-> 'HReal ':-> 'HReal ':-> 'HMeasure HBool) +chal1 = + lam $ \sigma -> + lam $ \a -> + lam $ \b -> + lam $ \c -> + normal a sigma >>= \ya -> + normal b sigma >>= \yb -> + normal c sigma >>= \yc -> + dirac (yb < ya && yc < ya) + +chal2 :: (ABT Term abt) => abt '[] + ('HProb ':-> 'HReal ':-> 'HReal ':-> 'HReal ':-> 'HMeasure 'HReal) +chal2 = + lam $ \sigma -> + lam $ \a -> + lam $ \b -> + lam $ \c -> + normal a sigma >>= \ya -> + normal b sigma >>= \yb -> + normal c sigma >>= \yc -> + dirac (if_ (yb < ya && yc < ya) one zero) + +chal3 :: (ABT Term abt) => abt '[] ('HProb ':-> 'HMeasure 'HReal) +chal3 = lam $ \sigma -> app3 (app chal2 sigma) zero zero zero + +--seismic :: (ABT Term abt) => abt '[] +-- (SE.HStation +-- ':-> HPair 'HReal (HPair 'HReal (HPair 'HProb 'HReal)) +-- ':-> HPair 'HReal (HPair 'HReal (HPair 'HReal 'HProb)) +-- ':-> 'HMeasure 'HProb) +--seismic = lam3 (\s e d -> dirac $ SE.densT s e d) + +easyHMM :: (ABT Term abt) => abt '[] + ('HMeasure (HPair (HPair 'HReal 'HReal) (HPair 'HProb 'HProb))) +easyHMM = + gamma (prob_ 3) one >>= \noiseT -> + gamma_1_1 >>= \noiseE -> + normal zero noiseT >>= \x1 -> + normal x1 noiseE >>= \m1 -> + normal x1 noiseT >>= \x2 -> + normal x2 noiseE >>= \m2 -> + dirac (pair (pair m1 m2) (pair noiseT noiseE)) diff --git a/haskell/Tests/Sample.hs b/haskell/Tests/Sample.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/Sample.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE DataKinds + , GADTs + , FlexibleContexts + #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Tests.Sample where + +import Prelude hiding ((+)) +import GHC.Word (Word32) +import qualified Data.Vector as V + +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Syntax.Prelude +import Language.Hakaru.Syntax.Value +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Sample + +import Tests.Models + +import qualified System.Random.MWC as MWC +import Test.HUnit + +seed :: V.Vector Word32 +seed = V.singleton 42 + +testMeasure :: String + -> Value ('HMeasure a) + -> Value 'HProb + -> Value a + -> Assertion +testMeasure p (VMeasure m) w v = do + g <- MWC.initialize seed + Just (v', w') <- m (VProb 1) g + assertEqual p v' v + assertEqual p w' w + +testEvaluate :: (ABT Term abt) + => String + -> abt '[] a + -> Value a + -> Assertion +testEvaluate p prog = assertEqual p (runEvaluate prog) + +normal01Value :: Value ('HMeasure 'HReal) +normal01Value = runEvaluate (triv normal_0_1) + +allTests :: Test +allTests = test + [ testMeasure "normal01" normal01Value (VProb 1) (VReal 0.35378756491616103) + , testEvaluate "1+1" (triv $ real_ 1 + real_ 1) (VReal 2) + ] diff --git a/haskell/Tests/Simplify.hs b/haskell/Tests/Simplify.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/Simplify.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE OverloadedStrings + , DataKinds + , FlexibleContexts + #-} + +module Tests.Simplify where + +import Prelude hiding ((>>=)) + +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.Sing +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.Prelude +import Language.Hakaru.Simplify + +import Test.HUnit +import Tests.TestTools + +v :: (ABT Term abt) => abt '[] ('HMeasure 'HNat) +v = var (Variable "x" 0 (SMeasure SNat)) + +freevar :: TrivialABT Term '[] ('HMeasure 'HNat) +freevar = v + +normal01T :: TrivialABT Term '[] ('HMeasure 'HReal) +normal01T = + syn (MeasureOp_ Normal + :$ syn (Literal_ (LReal (-2))) + :* syn (Literal_ (LProb 1)) + :* End) + +realpair :: TrivialABT Term '[] ('HMeasure (HPair 'HReal 'HReal)) +realpair = + ann_ (SMeasure $ sPair SReal SReal) + (dirac (pair (nat2real $ nat_ 1) (nat2real $ nat_ 2))) + +unifprob :: TrivialABT Term '[] ('HMeasure 'HProb) +unifprob = + uniform (real_ 1) (real_ 2) >>= \x -> + dirac (unsafeProb x) + +testSimplify + :: ( ABT Term abt + , Show (abt '[] a) + , Eq (abt '[] a) + ) + => String + -> abt '[] a + -> abt '[] a + -> Assertion +testSimplify nm x y = do + x' <- simplify x + assertEqual nm y x' + +allTests :: Test +allTests = test + [ testSimplify "freevar" freevar freevar + , testSimplify "normal01T" normal01T normal01T + , testSimplify "realpair" realpair realpair + , testSimplify "unifprob" unifprob unifprob + , testS "true" (triv $ ann_ (SMeasure sBool) (dirac true)) + ] diff --git a/haskell/Tests/TestSuite.hs b/haskell/Tests/TestSuite.hs --- a/haskell/Tests/TestSuite.hs +++ b/haskell/Tests/TestSuite.hs @@ -3,12 +3,14 @@ import System.Exit (exitFailure) import System.Environment (lookupEnv) -import qualified Tests.Parser as P -import qualified Tests.TypeCheck as TC -import qualified Tests.Simplify as S -import qualified Tests.Disintegrate as D -import qualified Tests.Sample as E -import qualified Tests.RoundTrip as RT +import qualified Tests.ASTTransforms as TR +import qualified Tests.Parser as P +import qualified Tests.Pretty as Pr +import qualified Tests.TypeCheck as TC +import qualified Tests.Simplify as S +import qualified Tests.Disintegrate as D +import qualified Tests.Sample as E +import qualified Tests.RoundTrip as RT import Test.HUnit @@ -26,15 +28,22 @@ allTests :: Maybe String -> Test allTests env = test [ TestLabel "Parser" P.allTests + , TestLabel "Pretty" Pr.allTests , TestLabel "TypeCheck" TC.allTests , TestLabel "Simplify" (simplifyTests S.allTests env) , TestLabel "Disintegrate" D.allTests , TestLabel "Evaluate" E.allTests , TestLabel "RoundTrip" (simplifyTests RT.allTests env) + , TestLabel "ASTTransforms" TR.allTests ] main :: IO () -main = do +main = mainWith (fmap Just . runTestTT) + +mainWith :: (Test -> IO (Maybe Counts)) -> IO () +mainWith run = do env <- lookupEnv "LOCAL_MAPLE" - Counts _ _ e f <- runTestTT (allTests env) - if (e>0) || (f>0) then exitFailure else return () + run (allTests env) >>= + maybe (return ()) (\(Counts _ _ e f) -> if (e>0) || (f>0) then exitFailure else return ()) + +-- maini = mainWith diff --git a/haskell/Tests/TestTools.hs b/haskell/Tests/TestTools.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/TestTools.hs @@ -0,0 +1,162 @@ +{-# LANGUAGE DeriveDataTypeable + , DataKinds + , RankNTypes + , GADTs + , PolyKinds + , FlexibleContexts #-} +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +module Tests.TestTools where + +import Language.Hakaru.Types.Sing +import Language.Hakaru.Parser.Parser +import Language.Hakaru.Parser.SymbolResolve (resolveAST) +import Language.Hakaru.Command (parseAndInfer, splitLines) +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.AST +import Language.Hakaru.Syntax.TypeCheck +import Language.Hakaru.Syntax.AST.Eq (alphaEq) +import Language.Hakaru.Syntax.IClasses (TypeEq(..), jmEq1) +import Language.Hakaru.Pretty.Concrete +import Language.Hakaru.Simplify +import Language.Hakaru.Syntax.AST.Eq() +import Text.PrettyPrint (Doc) + +import Data.Maybe (isJust) +import Data.List +import qualified Data.Text as T +import qualified Data.Text.IO as IO +import Data.Typeable (Typeable) +import Control.Exception +import Control.Monad + +import Test.HUnit + +data TestException = TestSimplifyException String SomeException + deriving Typeable +instance Exception TestException +instance Show TestException where + show (TestSimplifyException prettyHakaru e) = + show e ++ "\nwhile simplifying Hakaru:\n" ++ prettyHakaru + +-- assert that we get a result and that no error is thrown +assertResult :: [a] -> Assertion +assertResult s = assertBool "no result" $ not $ null s + +assertJust :: Maybe a -> Assertion +assertJust = assertBool "expected Just but got Nothing" . isJust + +handleException :: String -> SomeException -> IO a +handleException t e = throw (TestSimplifyException t e) + +testS + :: (ABT Term abt) + => String + -> abt '[] a + -> Assertion +testS p x = do + _ <- simplify x `catch` handleException (p ++ ": simplify failed") + return () + +testStriv + :: TrivialABT Term '[] a + -> Assertion +testStriv = testS "" + +-- Assert that all the given Hakaru programs simplify to the given one +testSS + :: (ABT Term abt) + => String + -> [(abt '[] a)] + -> abt '[] a + -> Assertion +testSS nm ts t' = + mapM_ (\t -> do p <- simplify t + assertAlphaEq nm p t') + (t':ts) + +testSStriv + :: [(TrivialABT Term '[] a)] + -> TrivialABT Term '[] a + -> Assertion +testSStriv = testSS "" + +assertAlphaEq :: + (ABT Term abt) + => String + -> abt '[] a + -> abt '[] a + -> Assertion +assertAlphaEq preface a b = + unless (alphaEq a b) (assertFailure $ mismatchMessage pretty preface a b) + +mismatchMessage :: forall (k :: q -> *) . (forall a . k a -> Doc) -> String -> forall a b . k a -> k b -> String +mismatchMessage k preface a b = msg + where msg = concat [ p + , "expected:\n" + , show (k b) + , "\nbut got:\n" + , show (k a) + ] + p = if null preface then "" else preface ++ "\n" + +testWithConcrete :: + (ABT Term abt) + => T.Text + -> TypeCheckMode + -> (forall a. Sing a -> abt '[] a -> Assertion) + -> Assertion +testWithConcrete s mode k = + case parseHakaru s of + Left err -> assertFailure (show err) + Right past -> + let m = inferType (resolveAST past) in + case runTCM m (splitLines s) mode of + Left err -> assertFailure (show err) + Right (TypedAST typ ast) -> k typ ast + + +testWithConcrete' + :: T.Text + -> TypeCheckMode + -> (forall a. Sing a -> TrivialABT Term '[] a -> Assertion) + -> Assertion +testWithConcrete' = testWithConcrete + +testConcreteFiles + :: FilePath + -> FilePath + -> Assertion +testConcreteFiles f1 f2 = do + t1 <- IO.readFile f1 + t2 <- IO.readFile f2 + case (parseAndInfer t1, parseAndInfer t2) of + (Left err, _) -> assertFailure (show err) + (_, Left err) -> assertFailure (show err) + (Right (TypedAST typ1 ast1), Right (TypedAST typ2 ast2)) -> do + ast1' <- simplify ast1 + ast2' <- simplify ast2 + case jmEq1 typ1 typ2 of + Just Refl -> assertAlphaEq "" ast1' ast2' + Nothing -> assertFailure "files don't have same type" + +ignore :: a -> Assertion +ignore _ = assertFailure "ignored" -- ignoring a test reports as a failure + +-- Runs a single test from a list of tests given its index +runTestI :: Test -> Int -> IO Counts +runTestI (TestList ts) i = runTestTT $ ts !! i +runTestI (TestCase _) _ = error "expecting a TestList, but got a TestCase" +runTestI (TestLabel _ _) _ = error "expecting a TestList, but got a TestLabel" + +hasLab :: String -> Test -> Bool +hasLab l (TestLabel lab _) = lab == l +hasLab _ _ = False + +-- Runs a single test from a TestList given its label +runTestN :: Test -> String -> IO Counts +runTestN (TestList ts) l = + case find (hasLab l) ts of + Just t -> runTestTT t + Nothing -> error $ "no test with label " ++ l +runTestN (TestCase _) _ = error "expecting a TestList, but got a TestCase" +runTestN (TestLabel _ _) _ = error "expecting a TestList, but got a TestLabel" diff --git a/haskell/Tests/TypeCheck.hs b/haskell/Tests/TypeCheck.hs new file mode 100644 --- /dev/null +++ b/haskell/Tests/TypeCheck.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE GADTs + , OverloadedStrings + , DataKinds + , FlexibleContexts + #-} + +module Tests.TypeCheck where + +import Prelude hiding (unlines) + +import qualified Language.Hakaru.Parser.AST as U +import qualified Language.Hakaru.Syntax.AST as T + +import Language.Hakaru.Syntax.AST.Eq() + +import Language.Hakaru.Syntax.ABT +import Language.Hakaru.Syntax.TypeCheck +import Language.Hakaru.Syntax.IClasses +import Language.Hakaru.Types.HClasses +import Language.Hakaru.Types.DataKind +import Language.Hakaru.Types.Sing + +import Data.Number.Nat + +import Data.Sequence +import Data.Text +import Test.HUnit +import Tests.TestTools + +five :: Text +five = "2 + 3" + +fiveU :: U.AST +fiveU = syn $ + U.NaryOp_ U.Sum + [ syn $ U.Literal_ $ Some1 $ T.LNat 2 + , syn $ U.Literal_ $ Some1 $ T.LNat 3 + ] + +fiveT :: TrivialABT T.Term '[] 'HNat +fiveT = + syn . T.NaryOp_ (T.Sum HSemiring_Nat) $ fromList + [ syn $ T.Literal_ $ T.LNat 2 + , syn $ T.Literal_ $ T.LNat 3 + ] + +normal01 :: U.AST +normal01 = syn $ + U.MeasureOp_ (U.SomeOp T.Normal) + [ syn $ U.Literal_ $ Some1 $ T.LReal 0 + , syn $ U.Literal_ $ Some1 $ T.LProb 1 + ] + +normal01T :: TrivialABT T.Term '[] ('HMeasure 'HReal) +normal01T = + syn (T.MeasureOp_ T.Normal + T.:$ (syn $ T.Literal_ $ T.LReal 0) + T.:* (syn $ T.Literal_ $ T.LProb 1) + T.:* T.End) + +xname :: Variable 'U.U +xname = Variable "x" (unsafeNat 0) U.SU + +normalb :: U.AST +normalb = syn $ + U.MBind_ + normal01 + (bind xname $ + syn $ U.MeasureOp_ (U.SomeOp T.Normal) + [ var xname + , syn $ U.Literal_ $ Some1 $ T.LProb 1 + ]) + + +inferType' :: U.AST -> TypeCheckMonad (TypedAST (TrivialABT T.Term)) +inferType' = inferType + +testTC :: Sing b -> U.AST -> TrivialABT T.Term '[] b -> Assertion +testTC typ uast tast = + case runTCM (inferType' uast) Nothing StrictMode of + Left _err -> assertFailure (show tast) + Right (TypedAST _typ ast) -> + case jmEq1 _typ typ of + Just Refl -> assertEqual "" tast ast + Nothing -> assertFailure + (show ast ++ " does not have same type as " ++ show tast) + +testConcreteTC :: Sing b -> Text -> TrivialABT T.Term '[] b -> Assertion +testConcreteTC typ s ast = + testWithConcrete' s StrictMode $ \_typ tast -> + case jmEq1 _typ typ of + Just Refl -> assertEqual "" tast ast + Nothing -> assertFailure + (show ast ++ " does not have same type as " ++ show tast) + + +allTests :: Test +allTests = test + [ testTC SNat fiveU fiveT + , testTC (SMeasure SReal) normal01 normal01T + , testConcreteTC SNat five fiveT + ]