{-# options_ghc -Wno-unused-imports -Wno-unused-top-binds #-}

module Core.Unification where

import Common
import Core.Config
import Core.Evaluation hiding
  (fresh, localVar, unfoldMeta, relevanceOf, relevanceInfoOf, cast, castFlex, castRigid
  , piApp, compareTypes, fltMax )
import Core.Evaluation qualified as Eval
import Core.Glued qualified as G
import Core.MetaState
import Core.Value
import Utils.UBool qualified as U
import Utils.StateT

{-

TODO:
- unify for levels
- lessThan for levels, with varying behavior depending on UnifyMode
- isSmall for types
- eta contraction to variables
- basic pattern unification (for now)

-}

-- UnifyM
----------------------------------------------------------------------------------------------------

type Res a = (# (# a, MetaState, U.Bool #) | (# MetaState, U.Bool #) #)

pattern OK :: a -> MetaState -> U.Bool -> Res a
pattern OK a s b <- (# (# a, s, b #) | #) where OK a s b = (# (# a, s, b #) | #)
{-# inline OK #-}

pattern Err :: MetaState -> U.Bool -> Res a
pattern Err s b = (# | (# s, b #) #)
{-# inline Err #-}
{-# complete OK, Err #-}

newtype UnifyM a = UnifyM# {runUnifyM :: MetaState -> U.Bool -> Res a}

pattern UnifyM :: (MetaState -> U.Bool -> Res a) -> UnifyM a
pattern UnifyM f <- UnifyM# f where
  UnifyM f = UnifyM# (oneShot \s -> oneShot \b -> f s b)
{-# inline UnifyM #-}
{-# complete UnifyM #-}

instance Functor UnifyM where
  {-# inline fmap #-}
  fmap = \f (UnifyM g) -> UnifyM \s b -> case g s b of
    OK a s b -> OK (f a) s b
    Err s b  -> Err s b
  {-# inline (<$) #-}
  (<$) = \a mb -> fmap (\_ -> a) mb

instance Applicative UnifyM where
  {-# inline pure #-}
  pure = \a -> UnifyM \s -> OK a s

  {-# inline (<*>) #-}
  (<*>) = \(UnifyM f) (UnifyM g) -> UnifyM \s b -> case f s b of
    OK f s b -> case g s b of
      OK a s b -> OK (f a) s b
      Err s b  -> Err s b
    Err s b -> Err s b

  {-# inline (<*) #-}
  (<*) = \(UnifyM f) (UnifyM g) -> UnifyM \s b -> case f s b of
    OK a s b -> case g s b of
      OK b s b' -> OK a s b'
      Err s b   -> Err s b
    e -> e

  {-# inline (*>) #-}
  (*>) = \(UnifyM f) (UnifyM g) -> UnifyM \s b -> case f s b of
    OK a s b -> case g s b of
      OK b s b' -> OK b s b'
      e         -> e
    Err s b -> Err s b

instance Monad UnifyM where
  {-# inline return #-}
  return = pure
  {-# inline (>>=) #-}
  (>>=) = \(UnifyM f) g -> UnifyM \s b -> case f s b of
    OK a s b -> runUnifyM (g a) s b
    Err s b  -> Err s b
  {-# inline (>>) #-}
  (>>) = (*>)

data UnifyState = US {
    unifyStateMetaState :: {-# nounpack #-} MetaState
  , unifyStateIsFlex :: Bool
  }

makeFields ''UnifyState

instance MonadState UnifyState UnifyM where
  {-# inline put #-}
  put = \(US s b) -> UnifyM \_ b -> OK () s b
  {-# inline get #-}
  get = UnifyM \s (U.Bool# b) -> OK (US s (isTrue# b)) s (U.Bool# b)

{-# inline withMetaState #-}
withMetaState :: (MetaStateArg => UnifyM a) -> UnifyM a
withMetaState k = get >>= \(US s _) -> let ?metaState = s in k

{-# inline liftMetaState #-}
liftMetaState :: (MetaStateArg => a) -> UnifyM a
liftMetaState k = withMetaState $ pure k

{-# inline localVar #-}
localVar :: ForceArgs => Lvl -> Lazy Ty -> UnifyM Tm
localVar l a = liftMetaState $ Eval.localVar l a

{-# inline fresh #-}
fresh :: LvlArg => Lazy Ty -> (LvlArg => Tm -> UnifyM a) -> UnifyM a
fresh a k = withMetaState $ Eval.fresh a k

{-# inline unfoldMeta #-}
unfoldMeta :: LvlArg => MetaHead -> a -> (Lazy Tm -> Bool -> UnifyM a) -> UnifyM a
unfoldMeta mh a k = withMetaState $ Eval.unfoldMeta mh (pure a) k

{-# inline relevanceInfoOf #-}
relevanceInfoOf :: LvlArg => Lazy Ty -> UnifyM (Lazy RelevanceInfo)
relevanceInfoOf a = liftMetaState $ Eval.relevanceInfoOf a

{-# inline relevanceOf #-}
relevanceOf :: LvlArg => Ty -> UnifyM Relevance
relevanceOf a = liftMetaState $ Eval.relevanceOf a

{-# inline piApp #-}
piApp :: LvlArg => Lazy Ty -> Lazy Tm -> UnifyM (Lazy Ty)
piApp a t = liftMetaState $ Eval.piApp a t

instance (c ~ LvlArg, a ~ Ty) => Apply PiClosure (Lazy Tm) (UnifyM a) c where
  (∙)  t u = liftMetaState $ t ∙ u

instance (c ~ LvlArg, a ~ Tm) => Apply LamClosure Tm (UnifyM a) c where
  (∙)  t u = liftMetaState $ t ∙ u

instance (c ~ LvlArg) => Apply Spine Tm (UnifyM Tm) c where
  (∙)  t u = liftMetaState $ t ∙ u

instance (c ~ LvlArg) => Apply Spine Tm (UnifyM (Lazy Tm)) c where
  (∙)  t u = liftMetaState $ Lazy (t ∙ u)

instance (c ~ LvlArg) => Apply Tm (Tm, PiDomKind) (UnifyM Tm) c where
  (∙)  t u = liftMetaState $ t ∙ u

instance (c ~ LvlArg) => Apply Tm (Tm, PiDomKind) (UnifyM (Lazy Tm)) c where
  (∙)  t u = liftMetaState $ Lazy (t ∙ u)

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

data UnifyMode
  = UMPure        -- ^ No metastate updates, no flex progress. Used in evaluation which
                  --   is completely pure.
  | UMDecision    -- ^ Used for deciding judgments without trying to make them become true.
                  --   Can add uniquely forced updates to the metastate. These are: solving
                  --   metas with contractible types, solving flex casts. The latter is forced
                  --   because every cast originates from an insertion into the elaboration output,
                  --   so every cast must be eventually solved.
                  --   We can't make flex progress.
  | UMImpure      -- ^ Updates the metastate in order to make judgements become true. Can solve
                  --   metas and make flex progress.
  deriving (Eq, Show)

type UnifyModeArg = (?unifyMode :: UnifyMode)
type UnifyArgs = (LvlArg, UnifyModeArg)

{-# inline putIsFlex #-}
putIsFlex :: U.Bool -> UnifyM ()
putIsFlex = \b -> UnifyM \s _ -> OK () s b

-- | Mark the current computation as flexbile. Throw error in non-Impure mode.
{-# inline markFlex #-}
markFlex :: UnifyModeArg => UnifyM ()
markFlex = UnifyM \s _ -> case ?unifyMode of
  UMImpure -> OK () s U.True
  _        -> Err s U.True

-- | Throw rigid mismatch
{-# inline throwRigid #-}
throwRigid :: UnifyM a
throwRigid = UnifyM Err

{-# inline guardMetaUpdate #-}
guardMetaUpdate :: UnifyModeArg => UnifyM ()
guardMetaUpdate = case ?unifyMode of
  UMImpure      -> pure ()
  _             -> throwRigid

{-# inline pureMode #-}
pureMode :: (UnifyModeArg => a) -> a
pureMode k = let ?unifyMode = UMPure in k

{-# inline impureMode #-}
impureMode :: (UnifyModeArg => a) -> a
impureMode k = let ?unifyMode = UMImpure in k

{-# inline decisionMode #-}
decisionMode :: (UnifyModeArg => a) -> a
decisionMode k = let ?unifyMode = UMDecision in k

{-# inline listenFlex #-}
-- | Run an action and check if it's flexible. Merge its flex-ness into the outer action.
listenFlex :: UnifyM a -> UnifyM (a, Bool)
listenFlex = \(UnifyM f) -> UnifyM \s b -> case f s U.False of
  OK a s (U.Bool# b') -> OK (a // isTrue# b') s (b U.|| U.Bool# b')
  Err s b -> Err s b

{-# inline listenFlex_ #-}
listenFlex_ :: UnifyM a -> UnifyM Bool
listenFlex_ = \act -> snd ! listenFlex act

{-# inline localFlex #-}
-- | Run an action but don't propagate its flex-ness into the outer action.
localFlex :: UnifyM a -> UnifyM a
localFlex = \(UnifyM f) -> UnifyM \s b -> case f s U.False of
  OK a s _ -> OK a s b
  Err s _  -> Err s b

-- | Branch on failure, keep the metastate changes.
{-# inline try #-}
try :: UnifyM a -> (a -> UnifyM b) -> UnifyM b -> UnifyM b
try = \(UnifyM f) succ (UnifyM fail) -> UnifyM \s b -> case f s b of
  OK a s b -> runUnifyM (succ a) s b
  Err s b  -> fail s b

-- | Run the second action on failure, keep all the metastate changes.
{-# inline retry #-}
retry :: UnifyM a -> UnifyM a -> UnifyM a
retry = \act act' -> try act pure act'

-- | Run the first action with meta solutions turned off, run the second one on failure.
{-# inline speculate #-}
speculate :: UnifyModeArg => (UnifyModeArg => UnifyM a) -> UnifyM a -> UnifyM a
speculate = \act act' ->
  (let ?unifyMode = case ?unifyMode of UMImpure -> UMDecision
                                       um       -> um
   in act)
  `retry` act'

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

class Unify a where
  unify :: UnifyArgs => a -> a -> UnifyM ()

unifyFlexDecide :: Unify a => UnifyArgs => a -> a -> UnifyM FlexDecide
unifyFlexDecide a a' =
  try (localFlex (listenFlex_ (unify a a')))
    (\case True  -> pure Dunno
           False -> pure Yes)
    (pure No)
{-# specialize unifyFlexDecide @Tm #-}
{-# specialize unifyFlexDecide @Ty #-}

{-# inline unifyEq #-}
unifyEq :: Eq a => UnifyModeArg => a -> a -> UnifyM ()
unifyEq = \a a' -> if a == a' then pure () else throwRigid

-- hs-boot exports
----------------------------------------------------------------------------------------------------
-- | Pure conversion.
convert :: ForceArgs => Unify a => a -> a -> FlexDecide
convert a a' =
  let ?unifyMode = UMPure in
  case runUnifyM (unify a a') ?metaState U.False of
    OK _ _ U.False -> Yes
    Err _  U.False -> No
    _              -> Dunno

-- | Monomorphized because it gets exported in hs-boot,
--   and hs-boot exports cannot be specialized or inlined.
convertTy :: ForceArgs => Ty -> Ty -> FlexDecide
convertTy = convert

-- | Monomorphized because it gets exported in hs-boot,
--   and hs-boot exports cannot be specialized or inlined
convertTm :: ForceArgs => Tm -> Tm -> FlexDecide
convertTm = convert

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

instance Unify PiClosure where
  unify b b' = do
    unifyEq (b^.domKind) (b'^.domKind)
    let a  = frc (b^.domTy)
        a' = frc (b'^.domTy)
    a2 <- listenFlex_ (unify a a')
    fresh (Lazy a) \x -> do
      cx   <- castIsFlex a2 (Lazy a) (Lazy a') x
      bx   <- b  ∙ Lazy x
      b'cx <- b' ∙ Lazy cx
      unify bx b'cx

instance Unify LamClosure where
  unify t t' = fresh (t^.domTy) \x -> do
    tx  <- t  ∙ x
    t'x <- t' ∙ x
    unify tx t'x

instance Unify a => Unify (Maybe a) where
  unify Nothing  Nothing   = pure ()
  unify (Just a) (Just a') = unify a a'
  unify _        _         = throwRigid

-- | Unify a Decode with a Decode. Requires its 'G.Tm' args to be forced!
decodeDecode :: UnifyArgs => Level -> G.Tm -> Level -> G.Tm -> UnifyM ()
decodeDecode s gt@(G.Tm t ft) s' gt'@(G.Tm t' ft') = case (ft, ft') of

  -- A rigid on either side forces the sizes
  (Rigid{}, _) -> do
    s2 <- listenFlex_ $ unify s s'
    ct <- castIsFlex s2 (Lazy (TYPE s)) (Lazy (TYPE s')) (frc t)
    unify ct (frc t')
  (_, Rigid{}) -> do
    s2  <- listenFlex_ $ unify s s'
    ct' <- castIsFlex s2 (Lazy (TYPE s')) (Lazy (TYPE s)) (frc t')
    unify (frc t) ct'

  -- speculate on unfold-unfold (TODO: heuristic)
  (Unfold h sp ft, Unfold h' sp' ft') ->
    (unify s s' >> unify h h' >> unify sp sp') `speculate` do
      ft  <- fhnf (frc ft)
      ft' <- fhnf (frc ft')
      decodeDecode s (G.Tm t ft) s' (G.Tm t' ft')

  -- lopsided unfold
  (Unfold _ _ ft, ft') -> do
    ft <- fhnf (frc ft)
    decodeDecode s (G.Tm t ft) s' gt'
  (ft, Unfold _ _ ft') -> do
    ft' <- fhnf (frc ft')
    decodeDecode s gt s' (G.Tm t' ft')

  -- remaining: flex-flex with flex distinct levels
  -- (assuming that we already passed the "shortcut" in
  -- Ty unification for decode-decode).
  _ -> markFlex

instance Unify Ty where
  unify a a' = unify (G.ty a) (G.ty a')

instance Unify G.Ty where
  unify (G.Ty topa ftopa) (G.Ty topa' ftopa') = do
    ftopa  <- fhnf ftopa
    ftopa' <- fhnf ftopa'
    case (ftopa, ftopa') of

      (Pi b    , Pi b'    ) -> unify b b'
      (FinLevel, FinLevel ) -> pure ()
      (Level   , Level    ) -> pure ()
      (TYPE s  , TYPE s'  ) -> unify s s'

      -- TODO shortcut: do a pure comparison of sizes
      --  if s = s', we can simply unify
      --     s < s', lift t
      --     s' < s, lift t'
      -- of course, the complication is that s,s' can be flex
      (Decode s t, Decode s' t') -> decodeDecode s (G.tm t) s' (G.tm t')

      (Decode _ (Flex FHCast{} _ _), _) -> markFlex
      (_, Decode _ (Flex FHCast{} _ _)) -> markFlex

      -- TODO: check smallness of canonical other side
      (Decode s (Flex (FHMeta mh) sp _), a') -> do
        uf
      (a, Decode s' (Flex (FHMeta mh') sp' _)) -> do
        uf

      _ -> throwRigid

instance Unify Tm where
  unify t t' = unify (G.tm t) (G.tm t')

instance Unify RigidHead where
  unify h h' = case (h, h') of
    (RHVar x, RHVar x') -> unifyEq x x'
    (RHCast a b t, RHCast a' b' t') -> do
      a2 <- listenFlex_ $ unify a a'
      unify b b'
      ct <- castIsFlex a2 (Lazy a) (Lazy a') t
      unify ct t'
    _ -> throwRigid

-- Spines
--------------------------------------------------------------------------------

unifyRSpine :: UnifyArgs => Bool -> RSpine -> RSpine -> UnifyM ()
unifyRSpine flex sp sp' = do
  sp  <- fhnf sp
  sp' <- fhnf sp'
  case (sp, sp') of
    (RSId, RSId) -> pure ()
    (RSApp t a pd sp, RSApp t' a' pd' sp') -> do
      ct <- castIsFlex flex a a' t
      t2 <- listenFlex_ $ unify ct t'
      unifyRSpine (flex || t2) sp sp'
    (RSCast (SCastFlex a b) sp, RSCast (SCastFlex a' b') sp') ->
      decisionMode do
        unify a a'
        unify b b'
        unifyRSpine False sp sp'
    (RSCast (SCastPi b1 b2) sp, RSCast (SCastPi b1' b2') sp') -> do
      unify b1 b1'
      b22 <- listenFlex_ $ unify b2 b2'
      unifyRSpine b22 sp sp'
    (RSCast (SCastRigid a b) sp, RSCast (SCastRigid a' b') sp') -> do
      unify a a'
      b2 <- listenFlex_ $ unify b b'
      unifyRSpine b2 sp sp'
    (RSCast SCastFlex{} _, _) -> markFlex
    (_, RSCast SCastFlex{} _) -> markFlex
    _ -> throwRigid


-- NOTE: these can be only called in
----------------------------------------------------------------------------------------------------

instance Unify UnfoldHead where
  unify h h' = case (h, h') of
    (UHMeta mh   , UHMeta mh'   ) -> unify mh mh'
    (UHTopDef i  , UHTopDef i'  ) -> if ptrEq i i' then pure () else throwRigid
    (UHLocalDef x, UHLocalDef x') -> unifyEq x x'
    _                             -> throwRigid

instance Unify MetaHead where
  unify (MetaHead m e) (MetaHead m' e') = unifyEq m m' >> unify e e'

instance Unify MetaEnv where
  unify e e' = unify (e^.env) (e'^.env)

instance Unify Env where
  unify e e' = case (e, e') of
    (ENil    , ENil      ) -> pure ()
    (EExt e t, EExt e' t') -> unify e e' >> unify (frc t) (frc t')
    (EDef e t, EDef e' t') -> unify e e' >> unify (frc t) (frc t')
    _                      -> throwRigid

instance Unify UnfoldSpine where
  unify sp sp'
    | unfoldSpineHasCast sp || unfoldSpineHasCast sp' =
      unify (reverseUnfoldSpine sp) (reverseUnfoldSpine sp')
    | otherwise =
      fastUnifyUnfoldSpine sp sp'

fastUnifyUnfoldSpine :: UnifyArgs => UnfoldSpine -> UnfoldSpine -> UnifyM ()
fastUnifyUnfoldSpine sp sp' = case (sp, sp') of
  (USId, USId) -> pure ()
  (USApp sp t pd, USApp sp' t' pd') -> do
    fastUnifyUnfoldSpine sp sp'
    unify t t'
  _ -> throwRigid

instance Unify RUnfoldSpine where
  unify sp sp' = do
    sp  <- fhnf sp
    sp' <- fhnf sp'
    case (sp, sp') of
      (RUSId, RUSId) -> pure ()
      (RUSApp t pd sp, RUSApp t' pd' sp') -> do
        unify t t'
        unify sp sp'
      (RUSCast (SCastFlex a b) sp, RUSCast (SCastFlex a' b') sp') ->
        decisionMode do
          unify a a'
          unify b b'
          unify sp sp'
      (RUSCast (SCastPi b1 b2) sp, RUSCast (SCastPi b1' b2') sp') -> do
        unify b1 b1'
        unify b2 b2'
        unify sp sp'
      (RUSCast (SCastRigid a b) sp, RUSCast (SCastRigid a' b') sp') -> do
        unify a a'
        b2 <- unify b b'
        unify sp sp'
      (RUSCast SCastFlex{} _, _) -> throwRigid

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

-- | Unify two meta-head Flex-s.
metaMeta :: UnifyArgs => MetaHead -> Spine -> Lazy Tm
                      -> MetaHead -> Spine -> Lazy Tm -> Lazy Ty -> UnifyM ()
metaMeta mh@(MetaHead m e) sp lhs mh'@(MetaHead m' e') sp' rhs a =
    -- We try to solve the newer metas first, to minimize the number of scope-escaping metas.a
  case compare m m' of
    LT -> listenFlex_ (solve mh' sp' lhs a) >>= \case
      True -> solve mh sp rhs a
      _    -> pure ()
    GT -> listenFlex_ (solve mh sp rhs a) >>= \case
      True -> solve mh' sp' lhs a
      _    -> pure ()
    -- TODO: intersect
    EQ -> (unify e e' >> unify sp sp') `speculate` markFlex

solve :: UnifyArgs => MetaHead -> Spine -> Lazy Tm -> Lazy Ty -> UnifyM ()
solve mh sp rhs a = do
  -- guardMetaUpdate
  uf

-- | Eta-expand based on the rhs, before trying to solve.
solveEtaLong :: UnifyArgs => MetaHead -> Spine -> Lazy Tm -> G.Tm -> Lazy Ty -> UnifyM ()
solveEtaLong mh sp lhs grhs@(G.Tm rhs frhs) a = whnf frhs >>= \case
  Lam t -> do
    let dom = t^.domTy
        pd = t^.domKind
    fresh dom \x -> do
      sp   <- pure $ SApp sp x dom pd
      lhs  <- frc lhs ∙ (x, pd)
      grhs <- G.Tm ! (frc rhs ∙ (x, pd)) ∘ (frhs ∙ (x, pd))
      a    <- piApp a (Lazy x)
      solveEtaLong mh sp lhs grhs a
  Flex (FHMeta mh') sp' _ ->
    metaMeta mh sp lhs mh' sp' rhs a
  _ ->
    solve mh sp rhs a

-- | Try to solve metavariable without eta-expansions. If it fails, we retry with expansions. NOTE:
--   we don't roll back the metastate on failure, because everything is stable under eta.
solveEtaShort :: UnifyArgs => MetaHead -> Spine -> Lazy Tm -> G.Tm -> Lazy Ty -> UnifyM ()
solveEtaShort mh sp lhs grhs@(G.Tm rhs _) a = do
  solve mh sp rhs a `retry` solveEtaLong mh sp lhs grhs a

-- | Unify spines when they don't contain any casts.
fastUnifySpine :: UnifyArgs => Spine -> Spine -> UnifyM ()
fastUnifySpine sp sp' = case (sp, sp') of
  (SId, SId) -> pure ()
  (SApp sp t a pd, SApp sp' t' a' pd') -> do
    sp2 <- listenFlex_ (fastUnifySpine sp sp')
    unifyEq pd pd'
    ct <- castIsFlex sp2 a a' t
    unify ct t'
  _ -> throwRigid

instance Unify Spine where
  unify sp sp'
    | spineHasCast sp || spineHasCast sp' =
      unifyRSpine False (reverseSpine sp) (reverseSpine sp')
    | otherwise =
      fastUnifySpine sp sp'

-- TODO note: more efficient if this /expects/ fnhf-d second component,
-- because then we don't re-force both sides on lopsided unfolding.
-- (but of course that's less defensive programming! Forcing everthing at
--  matching sites is the most defensive)
instance Unify G.Tm where
  unify (G.Tm topt ftopt) (G.Tm topt' ftopt') = do
    ftopt  <- fhnf ftopt
    ftopt' <- fhnf ftopt'
    case (ftopt, ftopt') of

      -- Match
      --------------------------------------------------------------------------------

      (Rigid h sp i, Rigid h' sp' _) -> do
        unify h h'
        s <- get
        try (unify sp sp')
          (\_ -> do
            s' <- get
            if sameMetaState (s^.metaState) (s'^.metaState) then
              pure ()
            else do
              whnf (frc i ^. relevance) >>= \case
                Relevant      -> pure ()
                Contractible  -> put s
                RelevanceOf{} -> put s >> markFlex
          )
          (do
            whnf (frc i ^. relevance) >>= \case
              Relevant      -> throwRigid
              Contractible  -> put s
              RelevanceOf{} -> put s >> markFlex
          )

      (Lam t, Lam t') -> unify t t'
      (Code _ a, Code _ a') -> unify a a'

      -- TODO: use unfolding heuristics (e.g. the Lean one)
      (Unfold h sp t, Unfold h' sp' t') ->
        (unify h h' >> unify sp sp')
         `speculate`
         unify (G.Tm topt (frc t)) (G.Tm topt' (frc t'))

      (FinLevelTm t, FinLevelTm t') -> unify t t'

      (Flex (FHMeta mh) sp a, Flex (FHMeta mh') sp' _) ->
        metaMeta mh sp topt mh' sp' topt' a

      (Flex (FHCast a b t) sp _, Flex (FHCast a' b' t') sp' ty) ->
        (do
          unify a a'
          unify t t'
          unify b b'
          unify sp sp'
         ) `speculate` markFlex

      -- Mismatch
      --------------------------------------------------------------------------------

      -- eta-short meta solutions
      (t@(Flex (FHMeta mh) sp a), t') ->
        solveEtaShort mh sp topt (G.Tm topt' t') a
      (t, t'@(Flex (FHMeta mh') sp' a')) ->
        solveEtaShort mh' sp' topt' (G.Tm topt t) a'

      -- lopsided unfolding
      (Unfold _ _ t, t') -> unify (G.Tm topt (frc t)) (G.Tm topt' t')
      (t, Unfold _ _ t') -> unify (G.Tm topt t) (G.Tm topt' (frc t'))

      -- syntax-directed eta
      (Lam t, t') -> fresh (t^.domTy) \x -> do
        let xpd = (x, t^.domKind)
        tx     <- t  ∙ x
        t'x    <- t' ∙ xpd
        toptx  <- frc topt  ∙ xpd
        topt'x <- frc topt' ∙ xpd
        unify (G.Tm toptx tx) (G.Tm topt'x t'x)
      (t, Lam t') -> fresh (t'^.domTy) \x -> do
        let xpd = (x, t'^.domKind)
        tx     <- t ∙ xpd
        t'x    <- t' ∙ x
        toptx  <- frc topt  ∙ xpd
        topt'x <- frc topt' ∙ xpd
        unify (G.Tm toptx tx) (G.Tm topt'x t'x)

      -- FHCast
      (Flex FHCast{} _ _, _) -> markFlex
      (_, Flex FHCast{} _ _) -> markFlex

      -- Handling identity lifting. The non-Code side must be Rigid at this point,
      -- since we already got rid of Unfold, flex meta and flex cast.
      (Code s1 (Decode s2 t), t') -> do
        seq <- listenFlex_ $ unify s1 s2
        ct <- castIsFlex seq (Lazy (TYPE s1)) (Lazy (TYPE s2)) t
        unify (G.tm t) (G.Tm topt' t')
      (t, Code s1' (Decode s2' t')) -> do
        seq <- listenFlex_ $ unify s1' s2'
        ct <- castIsFlex seq (Lazy (TYPE s1')) (Lazy (TYPE s2')) t
        unify (G.Tm topt t) (G.tm t')

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

      _ -> throwRigid


-- Impure casting
----------------------------------------------------------------------------------------------------

castFlex :: LvlArg => Ty -> Ty -> Tm -> UnifyM Tm
castFlex a b t = whnfImpure t >>= \case
  Rigid h sp _ -> liftMetaState $ castFlexRigid a b t h sp
  Flex h sp _  -> liftMetaState $ castFlexFlex a b t h sp
  _            -> pure $ Flex (FHCast a b t) SId (Lazy b)

castRigid :: LvlArg => Ty -> Ty -> Tm -> UnifyM Tm
castRigid a b t = whnfImpure t >>= \case
  Rigid h sp _ -> liftMetaState $ castRigidRigid a b t h sp
  Flex h sp _  -> liftMetaState $ castRigidFlex a b t h sp
  _            -> Rigid (RHCast a b t) SId ! relevanceInfoOf (Lazy b)

{-# inline impureCompareTypes #-}
impureCompareTypes :: LvlArg => Ty -> Ty -> UnifyM CompareTypes
impureCompareTypes a b = impureMode do
  a <- whnfImpure a
  b <- whnfImpure b
  case (a, b) of
    (Pi b, Pi b') -> pure $ CTPi b b'
    (a, b) -> unifyFlexDecide a b >>= \case
      Yes   -> pure CTSame
      No    -> pure CTDistinct
      Dunno -> pure CTFlex

-- Can be only called in Impure mode!!
cast :: LvlArg => Ty -> Ty -> Tm -> UnifyM Tm
cast a b t = do
  impureCompareTypes a b >>= \case
    CTPi b b'  -> liftMetaState $ castPi b b' t
    CTSame     -> pure t
    CTDistinct -> castRigid a b t
    CTFlex     -> castFlex a b t

-- Can be only called in Impure mode!!
{-# inline castIsFlex #-}
castIsFlex :: LvlArg => IsFlex -> Lazy Ty -> Lazy Ty -> Tm -> UnifyM Tm
castIsFlex flex a b t = case flex of
  True -> castFlex (frc a) (frc b) t
  _    -> pure t

-- Impure head forcing
----------------------------------------------------------------------------------------------------

instance Force RSpine where
  whnfImpure = \case
    sp@(RSCast cst sp') -> do
      let whnfCastPi :: LvlArg => RSpine -> PiClosure -> PiClosure -> RSpine -> UnifyM RSpine
          whnfCastPi sp b b' sp' = whnfImpure sp' >>= \case
            RSApp t a pd sp' -> do
              ct <- cast (frc (b'^.domTy)) (frc (b^.domTy)) t
              bct <- b  ∙ Lazy ct
              b't <- b' ∙ Lazy t
              pure $ RSApp ct (b^.domTy) (b^.domKind) $ RSCast (SCastFlex bct b't) sp'
            _ -> pure sp
      case cst of
        SCastFlex a b -> impureCompareTypes a b >>= \case
          CTPi b b'  -> whnfCastPi sp b b' sp'
          CTSame     -> whnfImpure sp'
          CTDistinct -> pure $ RSCast (SCastRigid a b) sp'
          CTFlex     -> pure sp
        SCastPi b b' -> whnfCastPi sp b b' sp'
        _ -> pure sp
    sp -> pure sp

  fhnfImpure = whnfImpure

instance Force RUnfoldSpine where
  whnfImpure = \case
    sp@(RUSCast cst sp') -> do
      let whnfCastPi :: LvlArg => RUnfoldSpine -> PiClosure -> PiClosure -> RUnfoldSpine
                     -> UnifyM RUnfoldSpine
          whnfCastPi sp b b' sp' = whnfImpure sp' >>= \case
            RUSApp t pd sp' -> do
              ct <- cast (frc (b'^.domTy)) (frc (b^.domTy)) t
              bct <- b  ∙ Lazy ct
              b't <- b' ∙ Lazy t
              pure $ RUSApp ct (b^.domKind) $ RUSCast (SCastFlex bct b't) sp'
            _ -> pure sp
      case cst of
        SCastFlex a b -> impureCompareTypes a b >>= \case
          CTPi b b'  -> whnfCastPi sp b b' sp'
          CTSame     -> whnfImpure sp'
          CTDistinct -> pure $ RUSCast (SCastRigid a b) sp'
          CTFlex     -> pure sp
        SCastPi b b' -> whnfCastPi sp b b' sp'
        _ -> pure sp
    sp -> pure sp

  fhnfImpure = whnfImpure

class Force a where
  -- | Compute fully to whnf while retrying flexible casts
  --   and updating the metastate.
  whnfImpure :: LvlArg => a -> UnifyM a

  -- | Only compute flexible heads, don't unfold defs and solved
  --   metas.
  fhnfImpure :: LvlArg => a -> UnifyM a

{-# inline fhnf #-}
fhnf :: UnifyArgs => Force a => ForcePure a => a -> UnifyM a
fhnf a = case ?unifyMode of
  UMPure -> liftMetaState $ fhnfPure a
  _      -> fhnfImpure a

{-# inline whnf #-}
whnf :: UnifyArgs => Force a => ForcePure a => a -> UnifyM a
whnf a = case ?unifyMode of
  UMPure -> liftMetaState $ whnfPure a
  _      -> whnfImpure a

instance Force Ty where
  whnfImpure = \case
    a@(Decode s t) -> do
      ft <- whnfImpure t
      if ptrEq t ft then pure a else whnfImpure $ decode s ft
    a -> pure a

  fhnfImpure = \case
    a@(Decode s t) -> do
      ft <- fhnfImpure t
      if ptrEq t ft then pure a else fhnfImpure $ decode s ft
    a -> pure a


-- Levels
----------------------------------------------------------------------------------------------------

-- Note: mostly a copypasta of Evaluation.fltMax
fltMax :: ForceArgs => FinLevelTm -> FinLevelTm -> UnifyM FinLevelTm
fltMax t t' = let
  forgetFlex (FLT _ n tail) = (n, tail)

  -- Insert a single new component into the tail.
  insertTm :: ForceArgs => Tm -> Integer -> FLTTail -> UnifyM FLTTail
  insertTm t n = \case
    FLTNil ->
      pure FLTNil
    FLTCons t' n' tail -> decisionMode (unifyFlexDecide t t') >>= \case
      Yes   -> pure $ FLTCons t' (max n n') tail
      No    -> FLTCons t' n' ! insertTm t n tail
      Dunno -> do putIsFlex U.True; FLTCons t' n' ! insertTm t n tail

  -- Insert an FLT into an FLT
  insertFlt :: ForceArgs => (Integer, FLTTail) -> (Integer, FLTTail) -> UnifyM (Integer, FLTTail)
  insertFlt (n, tail) (n', tail') = do
    (n'', tail'') <- insertList tail tail'
    pure (n `max` n' `max` n'' // tail'')

  -- Insert a tail into a tail. The Integer output is needed because
  -- a Tm can be forced to a FinLevelTm, which contains (Integer, FLTTail),
  -- and that has to be folded into the overall result.
  insertList :: ForceArgs => FLTTail -> FLTTail -> UnifyM (Integer, FLTTail)
  insertList tail tail' = case tail of
    FLTNil ->
      pure (0, tail')
    FLTCons t n tail -> whnfImpure t >>= \case
      FinLevelTm (forgetFlex -> (n', tail'')) ->
        insertFlt (n + n' // integerPlusFLTTail n tail'') =<< insertList tail tail'
      t@Flex{} -> do putIsFlex U.True; insertList tail =<< insertTm t n tail'
      t        -> insertList tail =<< insertTm t n tail'

  top :: UnifyM (Integer, FLTTail)
  top = insertFlt (forgetFlex t) =<< insertFlt (forgetFlex t') (forgetFlex fltZero)

  in localFlex (listenFlex top) >>= \case
   ((n, tail), flex) -> pure $ FLT flex n tail


instance Unify Level where
  unify l l' = case (l, l') of
    (ZeroPlus t    , ZeroPlus t'   ) -> unify t t'
    (OmegaPlus l   , OmegaPlus l'  ) -> unifyEq l l'
    (OmegaPlusOmega, OmegaPlusOmega) -> pure ()
    _                                -> throwRigid

instance Unify FinLevelTm where
  -- unify t t' = do
  --   FLT n tail <- whnf

instance Force FinLevelTm where
  whnfImpure = \case
    flt@FLTRigid{} -> pure flt
    flt@FLTFlex{}  -> uf -- fltMax flt fltZero

  fhnfImpure = whnfImpure

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

instance Force Tm where
  whnfImpure = \case
    t@(Flex h sp topa) -> case h of
      FHMeta mh     -> unfoldMeta mh t \v _ -> whnfImpure =<< (sp ∙ frc v)
      FHCast a b t' -> do ct' <- cast a b t'
                          whnfImpure =<< (sp ∙ ct')
    Unfold _ _ t -> whnfImpure (frc t)
    t            -> pure t

  fhnfImpure = \case
    t@(Flex h sp topa) -> case h of
      FHMeta mh     -> unfoldMeta mh t \v _ ->
                         Unfold (UHMeta mh) (spineToUnfoldSpine sp) ! (sp ∙ frc v)
      FHCast a b t' -> do ct' <- cast a b t'
                          fhnfImpure =<< (sp ∙ ct')
    t -> pure t

instance Force Relevance where
  whnfImpure = \case
    rel@(RelevanceOf s h sp) -> do
      let t = Flex h sp (Lazy (TYPE s))
      whnfImpure t >>= \case
        Code _ a         -> whnfImpure =<< relevanceOf a
        Rigid h sp _     -> pure Relevant
        t'@(Flex h sp _) -> if ptrEq t t' then pure rel else pure $ RelevanceOf s h sp
        _                -> impossible
    rel -> pure rel

  fhnfImpure = whnfImpure