EqSolver.hs

{-|
  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)
import Clash.Core.DataCon (dcUniq)
import GHC.Builtin.Names (unsafeReflDataConKey, eqPrimTyConKey, typeNatAddTyFamNameKey)

-- | 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
isAbsurdPat _tcm (DataPat dc _ _)
  -- unsafeCoerce is not absurd in the way intended by /isAbsurdPat/
  | dcUniq dc == fromGhcUnique unsafeReflDataConKey
  = False
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