Elaboration.hs

module Elaboration (elab) where

import GHC.Exts

import qualified LvlSet as LS
import qualified MetaCxt as MC
import qualified Presyntax as P
import qualified SymTable as ST
import qualified TopCxt as Top
import qualified Unification as Unif

import Common
import CoreTypes
import Cxt
import Exceptions
import InCxt

--------------------------------------------------------------------------------

unify :: Cxt -> P.Tm -> G -> G -> IO ()
unify cxt pt l r = do
  debug ["unify", showValOpt cxt (g1 l) UnfoldMetas, showValOpt cxt (g1 r) UnfoldMetas]
  let ecxt = ErrorCxt (mcxt cxt) (tbl cxt) (names cxt) (lvl cxt); {-# inline ecxt #-}
  Unif.unify (mcxt cxt) (lvl cxt) (frz cxt) Rigid l r `catch` \case
     UnifyEx e -> throw $ UnifyExInCxt ecxt pt (g1 l) (g1 r) e
     _         -> impossible

-- Fresh metas and meta insertions
--------------------------------------------------------------------------------

-- | Term, unforced type, forced type.
data Infer = Infer Tm {-# unpack #-} GTy

-- | Create fresh meta both as a term and a value.
freshMeta :: Cxt -> IO (Tm, Val)
freshMeta cxt = do

  let goSp x l mask acc
        | x == l    = acc
        | otherwise =
          let acc' | LS.member x mask = SApp acc (VLocalVar x SId) Expl
                   | otherwise        = acc
          in goSp (x + 1) l mask acc'

  mvar <- MC.fresh (mcxt cxt) (mask cxt)
  let mt = InsertedMeta (coerce mvar)
  let mv = VFlex (coerce mvar) (goSp 0 (lvl cxt) (mask cxt) SId)
  pure $! (mt // mv)
{-# inline freshMeta #-}

-- | Create fresh meta as a term, under an extra binder.
freshMetaUnderBinder :: Cxt -> Icit -> IO Tm
freshMetaUnderBinder cxt i = do
  mvar <- MC.fresh (mcxt cxt) (LS.insert (lvl cxt) (mask cxt))
  pure $! (InsertedMeta (coerce mvar))
{-# inline freshMetaUnderBinder #-}

goInsert' :: Cxt -> Infer -> IO Infer
goInsert' cxt (Infer t (G a fa)) = forceAll cxt fa >>= \case
  VPi (NI x Impl) a b -> do
    (m, mv) <- freshMeta cxt
    let b' = appCl' cxt b mv
    goInsert' cxt (Infer (App t m Impl) (gjoin b'))
  fa -> pure (Infer t (G a fa))

-- | Insert fresh implicit applications.
insertApps' :: Cxt -> IO Infer -> IO Infer
insertApps' cxt act = goInsert' cxt =<< act
{-# inline insertApps' #-}

-- | Insert fresh implicit applications to a term which is not
--   an implicit lambda (i.e. neutral).
insertApps :: Cxt -> IO Infer -> IO Infer
insertApps cxt act = act >>= \case
  res@(Infer (Lam (NI _ Impl) _) _) -> pure res
  res                               -> insertApps' cxt (pure res)
{-# inline insertApps #-}

-- | Insert fresh implicit applications until we hit a Pi with
--   a particular binder name.
insertAppsUntilName :: P.Tm -> Cxt -> Span -> IO Infer -> IO Infer
insertAppsUntilName topT cxt name act = go cxt =<< act where
  go cxt (Infer t (G topa ftopa)) = forceAll cxt ftopa >>= \case
    fa@(VPi (NI x Impl) a b) -> do
      if eqName cxt x (NSpan name) then
        pure (Infer t (G topa fa))
      else do
        (m, mv) <- freshMeta cxt
        let b' = appCl' cxt b mv
        go cxt (Infer (App t m Impl) (gjoin b'))
    _ ->
      throw $ NoNamedArgument topT name
{-# inline insertAppsUntilName #-}

-- Elaboration
--------------------------------------------------------------------------------

infer :: Cxt -> P.Tm -> IO Infer
infer cxt topT = do
  debug ["infer", showPTm cxt topT]

  Infer t a <- case topT of

    P.Var px -> do
      ma <- ST.lookup px (tbl cxt)

      case ma of
        UNothing             -> throw $ NotInScope px
        UJust (ST.Local x a) -> do

          -- debugging do
          --   foo <- io $ ST.assocs (tbl cxt)
          --   debug ["local var", show foo, showSpan (src cxt) px, show x,
          --          show $ lvlToIx (lvl cxt) x, show $ lvl cxt]

          pure $! (Infer (LocalVar (lvlToIx (lvl cxt) x)) a)

        UJust (ST.Top _ ga _ t) -> do

          -- debugging do
          --   foo <- io $ ST.assocs (tbl cxt)
          --   let x = case t of TopVar x _ -> x; _ -> impossible
          --   debug ["top var", show foo, showSpan (src cxt) px, show x]

          pure (Infer t ga)

    P.Let _ x ma t u ->
      checkWithAnnot cxt ma t \ t a va -> do
        Infer u vb <- defining cxt x va (eval cxt t) \cxt -> infer cxt u
        pure $! Infer (Let x a t u) vb

    topT@(P.App t u inf) -> do

      (i, t, G tty ftty) <- case inf of
        P.Named x -> do
          Infer t tty <- insertAppsUntilName topT cxt x $ infer cxt t
          pure (Impl, t, tty)
        P.NoName Impl -> do
          Infer t tty <- infer cxt t
          pure (Impl, t, tty)
        P.NoName Expl -> do
          Infer t tty  <- insertApps' cxt $ infer cxt t
          pure (Expl, t, tty)

      (a, b) <- forceAll cxt tty >>= \case
        VPi (NI x i') a b | i == i'   -> pure (a, b)
                          | otherwise -> impossible
        ftty -> do
          a <- snd <$> freshMeta cxt
          b <- Closure (env cxt) <$> freshMetaUnderBinder cxt i
          let expected = VPi (NI NX i) a b
          unify cxt topT (G tty ftty) (gjoin expected)
          pure (a, b)

      u <- check cxt u (gjoin a)
      let b' = appCl cxt b (eval cxt u)
      pure $! Infer (App t u i) (gjoin b')

    P.Pi _ x i a b -> do
      a <- checkType cxt a
      let ~va = eval cxt a
      binding cxt x i (gjoin va) \cxt _ -> do
        b <- checkType cxt b
        pure $! (Infer (Pi (NI (P.bind x) i) a b) (gjoin VU))

    P.Lam _ x P.Named{} ma t ->
      throw InferNamedLam

    P.Lam _ x (P.NoName i) ma t -> do
      (a, va) <- case ma of
        UNothing -> freshMeta cxt
        UJust a  -> do a <- checkType cxt a
                       pure (a, eval cxt a)

      Infer t vb <- binding cxt x i (gjoin va) \cxt _ -> insertApps cxt (infer cxt t)
      let ty = VPi (NI (P.bind x) i) va (valToClosure cxt (g1 vb))
      pure $! Infer (Lam (NI (P.bind x) i) t) (gjoin ty)

    P.U _ ->
      pure $! (Infer U (gjoin VU))

    P.Hole _ -> do
      (_, va) <- freshMeta cxt
      (t, _ ) <- freshMeta cxt
      pure $! (Infer t (gjoin va))

  debug ["inferred", showTm cxt t, showValOpt cxt (g1 a) UnfoldNone]
  pure (Infer t a)

-- | Choose between checking and inferring depending on an optional type annotation.
checkWithAnnot :: Cxt -> UMaybe P.Tm -> P.Tm -> (Tm -> Ty -> GTy -> IO a) -> IO a
checkWithAnnot cxt ma t k = case ma of
  UNothing -> do
    Infer t va <- insertApps cxt $ infer cxt t
    let a = quote cxt UnfoldNone (g1 va)
    k t a va
  UJust a -> do
    a <- checkType cxt a
    let va = gjoin $! eval cxt a
    t <- check cxt t va
    k t a va
{-# inline checkWithAnnot #-}

-- | Check that a preterm is a type.
checkType :: Cxt -> P.Tm -> IO Tm
checkType cxt pt = check cxt pt (gjoin VU)
{-# inline checkType #-}

check :: Cxt -> P.Tm -> GTy -> IO Tm
check cxt topT (G topA ftopA) = do
  debug ["check", showPTm cxt topT, showValOpt cxt topA UnfoldNone]

  ftopA <- forceAll cxt ftopA
  case (topT, ftopA) of
    (P.Lam _ x inf ma t, VPi (NI x' i) a b)
      | (case inf of P.NoName i' -> i == i'
                     P.Named n   -> eqName cxt x' (NSpan n) && i == Impl) -> do

      -- We prefer the user-provided annotation for the unforced binder type
      a <- case ma of
        UNothing -> pure a
        UJust a' -> do
          a' <- checkType cxt a'
          let va' = eval cxt a'
          unify cxt topT (gjoin va') (gjoin a)
          pure va'

      binding cxt x i (gjoin a) \cxt v -> do
        Lam (NI (P.bind x) i) <$> (check cxt t $! (gjoin $! appCl' cxt b v))

    (t, VPi (NI x Impl) a b) ->
      inserting cxt x \cxt v -> do
        t <- check cxt t $! (gjoin $! appCl' cxt b v)
        pure (Lam (NI x Impl) t)

    (P.Let _ x ma t u, ftopA) ->
      checkWithAnnot cxt ma t \ ~t ~a va ->
        defining cxt x va (eval cxt t) \cxt ->
          Let x a t <$> check cxt u (G topA ftopA)

    (P.Hole _, _) ->
      fst <$> freshMeta cxt

    (topT, ftopA) -> do
      Infer t inferred <- insertApps cxt $ infer cxt topT
      unify cxt topT inferred (G topA ftopA)
      pure t


-- Top level elaboration
--------------------------------------------------------------------------------

printingElabTime :: Top.WithCxt (Span -> Bool -> IO a -> IO a)
printingElabTime x timing act
  | timing = do
    (a, time) <- timed act
    putStrLn (showTopSpan x ++ " elaborated in " ++ show time)
    pure a
  | otherwise = act
{-# inline printingElabTime #-}

printingNfTime :: Top.WithCxt (Span -> Bool -> Tm -> IO ())
printingNfTime x timing t
  | timing = do
    let cxt = empty Top.cxt
    time <- timedPure_ (quote cxt UnfoldAll (eval cxt t))
    putStrLn (showTopSpan x ++ " normalized in " ++ show time)
  | otherwise = pure ()
{-# inline printingNfTime #-}

showTopSpan :: Top.WithCxt (Span -> String)
showTopSpan x = showSpan (ST.src ?tbl) x
{-# noinline showTopSpan #-}

elabTopLevel :: Top.WithCxt (P.TopLevel -> IO Top.Cxt)
elabTopLevel = \case

  P.Nil ->
    pure $! Top.cxt

  P.Definition (P.TopInfo x elabt nft) ma t u -> do

    let cxt = empty $! Top.cxt
    (t, a, va) <- case ma of

      UNothing -> do
        Infer t va <- printingElabTime x elabt (insertApps cxt $ infer cxt t)
        let a = quote cxt UnfoldNone (g1 va)
        printingNfTime x nft t
        pure (t, a, va)

      UJust a -> do
        a <- checkType cxt a
        let va = gjoin $! eval cxt a
        t <- printingElabTime x elabt (check cxt t va)
        printingNfTime x nft t
        pure (t, a, va)

    Top.define x a va t $ elabTopLevel u

elab :: Src -> P.TopLevel -> IO (Either Exception Top.Cxt)
elab src top =
  Top.new src (P.topLen top) (try $! elabTopLevel top)