module Elaborator where

import Data.List (elemIndex)
import Expr (Expr)
import qualified Expr as E
import IR (IRExpr)
import qualified IR as I

type Binders = [Text]

saveBinders :: State Binders a -> State Binders a
saveBinders action = do
    binders <- get
    res <- action
    put binders
    pure res

elaborate :: IRExpr -> Expr
elaborate ir = evalState (elaborate' ir) []
  where
    helper :: (Monad m) => Maybe a -> (a -> m b) -> m (Maybe b)
    helper Nothing _ = pure Nothing
    helper (Just x) f = Just <$> f x

    elaborate' :: IRExpr -> State Binders Expr
    elaborate' (I.Var n) = do
        binders <- get
        pure $ E.Var n . fromIntegral <$> elemIndex n binders ?: E.Free n
    elaborate' I.Star = pure E.Star
    elaborate' (I.Level level) = pure $ E.Level level
    elaborate' (I.App m n) = E.App <$> elaborate' m <*> elaborate' n
    elaborate' (I.Abs x t a b) = saveBinders $ do
        t' <- elaborate' t
        a' <- helper a elaborate'
        modify (x :)
        E.Abs x t' a' <$> elaborate' b
    elaborate' (I.Pi x t a b) = saveBinders $ do
        t' <- elaborate' t
        a' <- helper a elaborate'
        modify (x :)
        E.Pi x t' a' <$> elaborate' b
    elaborate' (I.Let name Nothing val body) = saveBinders $ do
        val' <- elaborate' val
        modify (name :)
        E.Let name Nothing val' <$> elaborate' body
    elaborate' (I.Let name (Just ty) val body) = saveBinders $ do
        val' <- elaborate' val
        ty' <- elaborate' ty
        modify (name :)
        E.Let name (Just ty') val' <$> elaborate' body