E.hs

type Lets = Exp -> Exp
-- "Stack of lets", functions of this form:
-- (Let b1 . Let b2 . ... . Let bn)

snoc :: Lets -> (Var, Exp) -> Lets
snoc k b = k . Let b

type M = State (T.Text, Int, Lets)

push :: (Var, Exp) -> M Exp
push b = modify' (\(s, c, k) -> (s, c, k `snoc` b))

wrapWithLet :: Exp -> M Exp
wrapWithLet exp = do
  (s, c, k) <- get
  let var = makeVar s c
  put (s, c+1, k `snoc` (var, exp))
  return $ Id var

flatten :: Exp -> M Exp
flatten (Sum e1 e2) = do
  exp1 <- flatten e1
  exp2 <- flatten e2
  wrapWithLet $ Sum exp1 exp2
flatten (Neg e) = do
  exp <- flatten e
  wrapWithLet $ Neg exp
flatten (Let (v, b) body) = do
  eBind <- flatten b
  push (v, eBind)
  flatten body
flatten e = return e

run :: Text -> M Exp -> Exp
run s r = let (i, (_, _, k)) = runState r (s, 0, id) in k i