-- | Equality solver.
--
-- Handles unification of meta-variables, decomposition of type
-- constructor equalities, and building coercion evidence.
module Aihc.Tc.Solve.Equality
  ( solveEquality,
    EqResult (..),
  )
where

import Aihc.Tc.Constraint
import Aihc.Tc.Evidence
import Aihc.Tc.Monad
import Aihc.Tc.Types
import Aihc.Tc.Zonk (zonkType)

-- | Result of attempting to solve an equality constraint.
data EqResult
  = -- | Solved: evidence bound.
    EqSolved
  | -- | Stuck: cannot solve yet (e.g. two different skolems).
    EqStuck !Ct
  | -- | Error: types are incompatible.
    EqError !Ct
  deriving (Show)

-- | Attempt to solve an equality constraint.
solveEquality :: Ct -> TcM EqResult
solveEquality ct = case ctPred ct of
  EqPred t1 t2 -> do
    t1' <- zonkType t1
    t2' <- zonkType t2
    solveEq (ct {ctPred = EqPred t1' t2'}) t1' t2'
  _ -> pure (EqStuck ct)

-- | Solve an equality between two zonked types.
solveEq :: Ct -> TcType -> TcType -> TcM EqResult
solveEq ct t1 t2 = case (t1, t2) of
  -- Same meta: trivially solved.
  (TcMetaTv u1, TcMetaTv u2) | u1 == u2 -> do
    bindEvidence (ctEvVar ct) (EvCoercion (Refl t1))
    pure EqSolved
  -- Meta on left: solve by binding.
  (TcMetaTv u, _) -> solveMetaEq ct u t2
  -- Meta on right: solve by binding.
  (_, TcMetaTv u) -> solveMetaEq ct u t1
  -- Same rigid variable.
  (TcTyVar v1, TcTyVar v2) | v1 == v2 -> do
    bindEvidence (ctEvVar ct) (EvCoercion (Refl t1))
    pure EqSolved
  -- Same type constructor: decompose.
  (TcTyCon tc1 args1, TcTyCon tc2 args2)
    | tc1 == tc2,
      length args1 == length args2 ->
        solveDecomposed ct t1 (zip args1 args2)
  -- Same function type shape: decompose.
  (TcFunTy a1 b1, TcFunTy a2 b2) ->
    solveDecomposed ct t1 [(a1, a2), (b1, b2)]
  -- Applied type constructor variable against a saturated tycon.
  (TcAppTy f a, TcTyCon tc args)
    | not (null args) ->
        solveDecomposed ct t1 [(f, TcTyCon tc (init args)), (a, last args)]
  (TcTyCon tc args, TcAppTy f a)
    | not (null args) ->
        solveDecomposed ct t1 [(TcTyCon tc (init args), f), (last args, a)]
  -- Same type-application shape. This is also needed when solving wanted
  -- equalities inside an implication, where canonicalization does not first
  -- decompose the application for us.
  (TcAppTy f1 a1, TcAppTy f2 a2) ->
    solveDecomposed ct t1 [(f1, f2), (a1, a2)]
  -- Incompatible types.
  _ -> pure (EqError ct)

-- | Solve a meta-variable equality by binding.
solveMetaEq :: Ct -> Unique -> TcType -> TcM EqResult
solveMetaEq ct u ty
  | occursIn u ty = pure (EqError ct)
  | otherwise = do
      writeMetaTv u ty
      bindEvidence (ctEvVar ct) (EvCoercion (Refl ty))
      pure EqSolved

solveDecomposed :: Ct -> TcType -> [(TcType, TcType)] -> TcM EqResult
solveDecomposed ct witness pairs = do
  results <- mapM solvePair pairs
  case firstUnsolved results of
    Nothing -> do
      bindEvidence (ctEvVar ct) (EvCoercion (Refl witness))
      pure EqSolved
    Just result -> pure result
  where
    solvePair (left, right) =
      solveEquality (ct {ctPred = EqPred left right})

firstUnsolved :: [EqResult] -> Maybe EqResult
firstUnsolved [] = Nothing
firstUnsolved (EqSolved : rest) = firstUnsolved rest
firstUnsolved (result : _) = Just result

-- | Occurs check: does meta-variable u appear in the type?
occursIn :: Unique -> TcType -> Bool
occursIn u = go
  where
    go (TcMetaTv u') = u == u'
    go (TcTyVar _) = False
    go (TcTyCon _ args) = any go args
    go (TcFunTy a b) = go a || go b
    go (TcForAllTy _ body) = go body
    go (TcQualTy preds body) = any goPred preds || go body
    go (TcAppTy f a) = go f || go a
    go (TcBuiltinTyCon _ _ arguments) = any go arguments

    goPred (ClassPred _ args) = any go args
    goPred (EqPred a b) = go a || go b
