content: {-|
    Copyright  :  (C) 2021-2024 QBayLogic B.V.
    License    :  BSD2 (see the file LICENSE)
    Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>
  -}
  
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Clash.Core.EqSolver where
  
  import Data.List.Extra (zipEqual)
  import Data.Maybe (catMaybes, mapMaybe)
  
  import Clash.Core.Name (Name(nameUniq))
  import Clash.Core.Term
  import Clash.Core.TyCon
  import Clash.Core.Type
  import Clash.Core.Var
  import Clash.Core.VarEnv (VarSet, elemVarSet, emptyVarSet, mkVarSet)
  import Clash.Unique (fromGhcUnique)
  #if MIN_VERSION_ghc(9,0,0)
  import Clash.Core.DataCon (dcUniq)
  import GHC.Builtin.Names (unsafeReflDataConKey, eqPrimTyConKey, typeNatAddTyFamNameKey)
  #else
  import PrelNames (eqPrimTyConKey, typeNatAddTyFamNameKey)
  #endif
  
  -- | Data type that indicates what kind of solution (if any) was found
  data TypeEqSolution
    = Solution (TyVar, Type)
    -- ^ Solution was found. Variable equals some integer.
    | AbsurdSolution
    -- ^ A solution was found, but it involved negative naturals.
    | NoSolution
    -- ^ Given type wasn't an equation, or it was unsolvable.
      deriving (Show, Eq)
  
  catSolutions :: [TypeEqSolution] -> [(TyVar, Type)]
  catSolutions = mapMaybe getSol
   where
    getSol (Solution s) = Just s
    getSol _ = Nothing
  
  -- | Solve given equations and return all non-absurd solutions
  solveNonAbsurds :: TyConMap -> VarSet -> [(Type, Type)] -> [(TyVar, Type)]
  solveNonAbsurds _tcm _ [] = []
  solveNonAbsurds tcm solveSet (eq:eqs) =
    solved ++ solveNonAbsurds tcm solveSet eqs
   where
    solvers = [pure . solveAdd solveSet, solveEq tcm solveSet]
    solved = catSolutions (concat [s eq | s <- solvers])
  
  -- | Solve simple equalities such as:
  --
  --   * a ~ 3
  --   * 3 ~ a
  --   * SomeType a b ~ SomeType 3 5
  --   * SomeType 3 5 ~ SomeType a b
  --   * SomeType a 5 ~ SomeType 3 b
  --
  solveEq :: TyConMap -> VarSet -> (Type, Type) -> [TypeEqSolution]
  solveEq tcm solveSet (coreView tcm -> left, coreView tcm -> right) =
    case (left, right) of
      (VarTy tyVar, ConstTy {}) | elemVarSet tyVar solveSet ->
        -- a ~ 3
        [Solution (tyVar, right)]
      (ConstTy {}, VarTy tyVar) | elemVarSet tyVar solveSet ->
        -- 3 ~ a
        [Solution (tyVar, left)]
      (ConstTy {}, ConstTy {}) ->
        -- Int /= Char
        if left /= right then [AbsurdSolution] else []
      (LitTy {}, LitTy {}) ->
        -- 3 /= 5
        if left /= right then [AbsurdSolution] else []
      _ ->
        -- The call to 'coreView' at the start of 'solveEq' should have reduced
        -- all solvable type families. If we encounter one here that means the
        -- type family is stuck (and that we shouldn't compare it to anything!).
        if any (isTypeFamilyApplication tcm) [left, right] then
          []
        else
          case (tyView left, tyView right) of
            (TyConApp leftNm leftTys, TyConApp rightNm rightTys) ->
              -- SomeType a b ~ SomeType 3 5 (or other way around)
              if leftNm == rightNm then
                concat (map (solveEq tcm solveSet) (zipEqual leftTys rightTys))
              else
                [AbsurdSolution]
            _ ->
              []
  
  -- | Solve equations supported by @normalizeAdd@. See documentation of
  -- @TypeEqSolution@ to understand the return value.
  solveAdd
    :: VarSet
    -> (Type, Type)
    -> TypeEqSolution
  solveAdd solveSet ab =
    case normalizeAdd ab of
      Just (n, m, VarTy tyVar) | elemVarSet tyVar solveSet ->
        if n >= 0 && m >= 0 && n - m >= 0 then
          Solution (tyVar, (LitTy (NumTy (n - m))))
        else
          AbsurdSolution
      _ ->
        NoSolution
  
  -- | Given the left and right side of an equation, normalize it such that
  -- equations of the following forms:
  --
  --     * 5     ~ n + 2
  --     * 5     ~ 2 + n
  --     * n + 2 ~ 5
  --     * 2 + n ~ 5
  --
  -- are returned as (5, 2, n)
  normalizeAdd
    :: (Type, Type)
    -> Maybe (Integer, Integer, Type)
  normalizeAdd (a, b) = do
    (n, rhs) <- lhsLit a b
    case tyView rhs of
      TyConApp tc [left, right]
        | nameUniq tc == fromGhcUnique typeNatAddTyFamNameKey -> do
        (m, o) <- lhsLit left right
        return (n, m, o)
      _ ->
        Nothing
   where
    lhsLit x                 (LitTy (NumTy n)) = Just (n, x)
    lhsLit (LitTy (NumTy n)) y                 = Just (n, y)
    lhsLit _                 _                 = Nothing
  
  -- | Tests for nonsencical patterns due to types being "absurd". See
  -- @isAbsurdEq@ for more info.
  isAbsurdPat
    :: TyConMap
    -> Pat
    -> Bool
  #if MIN_VERSION_base(4,15,0)
  isAbsurdPat _tcm (DataPat dc _ _)
    -- unsafeCoerce is not absurd in the way intended by /isAbsurdPat/
    | dcUniq dc == fromGhcUnique unsafeReflDataConKey
    = False
  #endif
  isAbsurdPat tcm pat =
    any (isAbsurdEq tcm exts) (patEqs tcm pat)
   where
    exts = case pat of
      DataPat _dc extNms _ids -> mkVarSet extNms
      _ -> emptyVarSet
  
  -- | Determines if an "equation" obtained through @patEqs@ or @typeEq@ is
  -- absurd. That is, it tests if two types that are definitely not equal are
  -- asserted to be equal OR if the computation of the types yield some absurd
  -- (intermediate) result such as -1.
  isAbsurdEq
    :: TyConMap
    -> VarSet -- ^ existential tvs
    -> (Type, Type)
    -> Bool
  isAbsurdEq tcm exts ((left0, right0)) =
    case (coreView tcm left0, coreView tcm right0) of
      (solveAdd exts -> AbsurdSolution) -> True
      lr -> any (==AbsurdSolution) (solveEq tcm exts lr)
  
  -- | Get constraint equations
  patEqs
    :: TyConMap
    -> Pat
    -> [(Type, Type)]
  patEqs tcm pat =
   catMaybes (map (typeEq tcm . varType) (snd (patIds pat)))
  
  -- | If type is an equation, return LHS and RHS.
  typeEq
    :: TyConMap
    -> Type
    -> Maybe (Type, Type)
  typeEq tcm ty =
   case tyView (coreView tcm ty) of
    TyConApp tc [_, _, left, right]
      | nameUniq tc == fromGhcUnique eqPrimTyConKey ->
      Just (coreView tcm left, coreView tcm right)
    _ ->
      Nothing
  