expr.hs

{-# LANGUAGE RecursiveDo #-}
module Expr where
  
import Control.Monad.Trans.Reader
import Data.Map (Map)
import qualified Data.Map as Map

type Var = String

data Expr
  = Add Expr Expr
  | Lit Int
  | Var Var
  | Let Var Expr Expr
  deriving Show

data Env a = Env { bindings :: Map Var a, recursiveBindings :: Map Var (Env a, Expr, Int) }

summarise :: Monoid a => (Int -> a) -> Expr -> a
summarise f e = runReader (go e) (Env Map.empty Map.empty)
  where
    go (Add a b) = mappend <$> go a <*> go b
    go (Lit i) = pure $ f i
    go (Var v) = do
      binds <- asks bindings
      pure (binds Map.! v)
    go (Let v e' e) = do
      a <- go e'
      local (\e -> e {bindings = Map.insert v a (bindings e) })
            (go e)

type UnrollLimit = Int

-- Unroll recursive let up to limit, then use `f 0`
summariseRec :: Monoid a => UnrollLimit -> (Int -> a) -> Expr -> a
summariseRec lim f e = runReader (go e) (Env Map.empty Map.empty)
  where
    go (Add a b) = mappend <$> go a <*> go b
    go (Lit i) = pure $ f i
    go (Var v) = do
      binds <- asks bindings
      case Map.lookup v binds of
        Just x -> pure x
        Nothing -> do
          recs <- asks recursiveBindings
          let (env,expr,i) = recs Map.! v
          if i < lim
          then local (\_ -> env { recursiveBindings = Map.insert v (env,expr,i+1) (recursiveBindings env)})
                     (go expr)
          else pure $ f 0
    go (Let v e' e) = mdo
      a <- local (\env -> let env' = env {recursiveBindings = Map.insert v (env',e',0) (recursiveBindings env) } in env')
                 (go e')
      local (\e -> e {bindings = Map.insert v a (bindings e) })
            (go e)