packages feed

futhark-0.27.1: src/Language/Futhark/TypeChecker/TySolve.hs

module Language.Futhark.TypeChecker.TySolve
  ( Type,
    Solution,
    UnconTyVar,
    solve,
  )
where

import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.ST
import Data.Bifunctor
import Data.List qualified as L
import Data.Loc
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.Util.Pretty
import Language.Futhark
import Language.Futhark.TypeChecker.Constraints
import Language.Futhark.TypeChecker.Error
import Language.Futhark.TypeChecker.Monad (Notes, TypeError (..), aNote, withIndexLink)
import Language.Futhark.TypeChecker.UnionFind

-- | The type representation used by the constraint solver. Agnostic
-- to sizes and uniqueness.
type Type = CtType ()

type UF s = M.Map TyVar (TyVarNode s)

newtype SolverState s = SolverState {solverTyVars :: UF s}

newtype SolveM s a = SolveM
  { runSolveM :: ExceptT TypeError (ReaderT (SolverState s) (ST s)) a
  }
  deriving (Functor, Applicative, Monad, MonadError TypeError, MonadReader (SolverState s))

-- | A solution maps a type variable to its substitution. This
-- substitution is complete, in the sense there are no right-hand
-- sides that contain a type variable.
type Solution = M.Map TyVar (Either [PrimType] (TypeBase () NoUniqueness))

-- | An unconstrained type variable comprises a name and (ironically)
-- a constraint on how it can be instantiated.
type UnconTyVar = (VName, Liftedness)

liftST :: ST s a -> SolveM s a
liftST = SolveM . lift . lift

getSol' :: TyVarNode s -> SolveM s TyVarSol
getSol' = liftST . getSol

union' :: TyVarNode s -> TyVarNode s -> SolveM s ()
union' tv1 tv2 = liftST $ union tv1 tv2

unionNewSol' tv1 tv2 new_sol = liftST $ unionNewSol tv1 tv2 new_sol

unionNewSol' :: TyVarNode s -> TyVarNode s -> TyVarSol -> SolveM s ()
getKey' :: TyVarNode s -> SolveM s TyVar
getKey' = liftST . getKey

initializeState :: TyParams -> TyVars () -> ST s (SolverState s)
initializeState typarams tyvars = do
  tyvars' <- M.traverseWithKey f tyvars
  typarams' <- M.traverseWithKey g typarams
  pure $ SolverState $ typarams' <> tyvars'
  where
    f tv (_lvl, info) = makeTyVarNode tv info
    g tv (lvl, lft, loc) = makeTyParamNode tv lvl lft loc

typeError :: Loc -> Notes -> Doc () -> SolveM s ()
typeError loc notes msg =
  throwError $ TypeError loc notes msg

typeVar :: (Monoid u) => VName -> TypeBase dim u
typeVar v = Scalar $ TypeVar mempty (qualName v) []

cannotUnify ::
  Reason Type ->
  Notes ->
  BreadCrumbs ->
  Type ->
  Type ->
  SolveM s ()
cannotUnify reason notes bcs t1 t2 = do
  t1' <- substTyVars t1
  t2' <- substTyVars t2
  case reason of
    Reason loc ->
      typeError loc notes . stack $
        [ "Cannot unify",
          indent 2 (pretty t1'),
          "with",
          indent 2 (pretty t2')
        ]
          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]
    ReasonPatMatch loc pat value_t ->
      typeError loc notes . stack $
        [ "Pattern",
          indent 2 $ align $ pretty pat,
          "cannot match value of type",
          indent 2 $ align $ pretty value_t
        ]
          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]
    ReasonAscription loc expected actual ->
      typeError loc notes . stack $
        [ "Expression does not have expected type from type ascription.",
          "Expected:" <+> align (pretty expected),
          "Actual:  " <+> align (pretty actual)
        ]
          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]
    ReasonRetType loc expected actual -> do
      expected' <- substTyVars expected
      actual' <- substTyVars actual
      typeError loc notes . stack $
        [ "Function body does not have expected type.",
          "Expected:" <+> align (pretty expected'),
          "Actual:  " <+> align (pretty actual')
        ]
          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]
    ReasonApply loc f e expected actual -> do
      expected' <- substTyVars expected
      actual' <- substTyVars actual
      typeError loc notes . stack $
        [ header,
          "Expected:" <+> align (pretty expected'),
          "Actual:  " <+> align (pretty actual')
        ]
      where
        header =
          case f of
            (Nothing, _) ->
              "Cannot apply function to"
                <+> dquotes (shorten $ group $ pretty e)
                <> " (invalid type)."
            (Just fname, _) ->
              "Cannot apply"
                <+> dquotes (pretty fname)
                <+> "to"
                <+> dquotes (align $ shorten $ group $ pretty e)
                <> " (invalid type)."
    ReasonApplySplit loc (fname, 0) _ ftype ->
      typeError loc notes $
        stack
          [ "Cannot apply"
              <+> fname'
              <+> "as function, as it has non-function type:"
              </> indent 2 (align $ pretty ftype)
          ]
      where
        fname' = maybe "expression" (dquotes . pretty) fname
    ReasonApplySplit loc (fname, i) e _ ->
      typeError loc notes $
        stack
          [ "Cannot apply"
              <+> fname'
              <+> "to"
              <+> dquotes (align $ shorten $ group $ pretty e)
              <> ".",
            "Function accepts only" <+> pretty i <+> "arguments."
          ]
      where
        fname' = maybe "expression" (dquotes . pretty) fname
    ReasonBranches loc former latter -> do
      former' <- substTyVars former
      latter' <- substTyVars latter
      typeError loc notes . stack $
        [ "Branches differ in type.",
          "Former:" <+> pretty former',
          "Latter:" <+> pretty latter'
        ]

unsharedConstructorsMsg :: M.Map Name t -> M.Map Name t -> Doc a
unsharedConstructorsMsg cs1 cs2 =
  "Unshared constructors:" <+> commasep (map (("#" <>) . pretty) missing) <> "."
  where
    missing =
      filter (`notElem` M.keys cs1) (M.keys cs2)
        ++ filter (`notElem` M.keys cs2) (M.keys cs1)

substTyVars :: (Monoid u) => TypeBase () u -> SolveM s (TypeBase () u)
substTyVars (Scalar (TypeVar u qn args)) = do
  mb_node <- maybeLookupUF $ qualLeaf qn
  case mb_node of
    Just node -> do
      sol <- getSol' node
      qn_k <- qualName <$> getKey' node
      case sol of
        Solved t -> do
          t' <- substTyVars t
          pure $ second (const mempty) t'
        _ -> makeTyVar qn_k
    _ -> makeTyVar qn
  where
    makeTyVar qn' = do
      args' <- mapM onArg args
      pure $ Scalar $ TypeVar u qn' args'
    onArg (TypeArgType t) = TypeArgType <$> substTyVars t
    onArg d@(TypeArgDim _) = pure d
substTyVars p@(Scalar (Prim _)) = pure p
substTyVars (Scalar (Record fs)) =
  Scalar . Record <$> traverse substTyVars fs
substTyVars (Scalar (Sum cs)) =
  Scalar . Sum <$> traverse (mapM substTyVars) cs
substTyVars (Scalar (Arrow u pname d t1 (RetType ext t2))) = do
  t1' <- substTyVars t1
  t2' <- substTyVars t2
  pure $
    Scalar $
      Arrow u pname d t1' $
        RetType ext $
          t2' `setUniqueness` uniqueness t2
substTyVars (Array u shape elemt) = do
  elemt' <- substTyVars $ Scalar elemt
  pure $ arrayOfWithAliases u shape elemt'

occursCheck :: Reason Type -> VName -> VName -> Type -> SolveM s ()
occursCheck reason v k tp = do
  let vars = typeVars tp
  when (k `S.member` vars)
    . typeError (locOf reason) mempty
    . withIndexLink "occurs-check"
    $ "Occurs check: cannot instantiate"
      <+> prettyName v
      <+> "with"
      <+> pretty tp
      <> "."

bindTyVar ::
  Reason Type ->
  BreadCrumbs ->
  VName ->
  TyVarNode s ->
  Type ->
  SolveM s ()
bindTyVar reason bcs v v_node t' = do
  t <- substTyVars t'
  k <- getKey' v_node
  occursCheck reason v k t

  v_info <- getSol' v_node

  setInfo v_node $ Solved t

  case (v_info, t) of
    (Unsolved TyVarFree {}, _) -> pure ()
    (Unsolved (TyVarPrim _ v_pts), _) ->
      if t `elem` map (Scalar . Prim) v_pts
        then pure ()
        else cannotUnify reason notes bcs (typeVar v) t
      where
        notes =
          aNote $
            "Cannot instantiate type that must be one of"
              </> indent 2 (pretty v_pts)
              </> "with"
              </> indent 2 (pretty t)
    (Unsolved (TyVarSum _ cs1), Scalar (Sum cs2)) ->
      if all (`elem` M.keys cs2) (M.keys cs1)
        then unifySharedConstructors reason bcs cs1 cs2
        else cannotUnify reason notes bcs (typeVar v) t
      where
        notes =
          aNote $
            "Cannot match type with constructors"
              </> indent 2 (stack (map (("#" <>) . pretty) (M.keys cs1)))
              </> "with type with constructors"
              </> indent 2 (stack (map (("#" <>) . pretty) (M.keys cs2)))
              </> unsharedConstructorsMsg cs1 cs2
    (Unsolved (TyVarSum _ cs1), _) ->
      typeError (locOf reason) mempty $
        "Cannot unify type with constructors"
          </> indent 2 (pretty (Sum cs1))
          </> "with type"
          </> indent 2 (pretty t)
    (Unsolved (TyVarRecord _ fs1), Scalar (Record fs2)) ->
      if all (`elem` M.keys fs2) (M.keys fs1)
        then unifySharedFields reason bcs fs1 fs2
        else
          typeError (locOf reason) mempty $
            "Cannot unify record type with fields"
              </> indent 2 (pretty (Record fs1))
              </> "with record type"
              </> indent 2 (pretty (Record fs2))
    (Unsolved (TyVarRecord _ fs1), _) ->
      typeError (locOf reason) mempty $
        "Cannot unify record type with fields"
          </> indent 2 (pretty (Record fs1))
          </> "with type"
          </> indent 2 (pretty t)
    --
    -- Internal error cases
    (Solved {}, _) ->
      error $ "Type variable already solved: " <> prettyNameString v
    (Param {}, _) ->
      error $ "Cannot substitute type parameter: " <> prettyNameString v

solveCt :: CtTy () -> SolveM s ()
solveCt (CtEq reason t1 t2) = solveEq reason mempty t1 t2

solveEq :: Reason Type -> BreadCrumbs -> Type -> Type -> SolveM s ()
solveEq reason obcs orig_t1 orig_t2 = do
  solveCt' (obcs, (orig_t1, orig_t2))
  where
    flexible :: VName -> SolveM s (Maybe (TyVarNode s))
    flexible v = do
      uf <- asks solverTyVars
      case M.lookup v uf of
        j_n@(Just node) -> do
          sol <- getSol' node
          pure $ case sol of
            Unsolved _ -> j_n
            _ -> Nothing
        Nothing -> pure Nothing

    normalize :: TypeBase () NoUniqueness -> SolveM s (TypeBase () NoUniqueness)
    normalize t@(Scalar (TypeVar _ (QualName [] v) [])) = do
      uf <- asks solverTyVars
      case M.lookup v uf of
        Just node -> do
          sol <- getSol' node
          case sol of
            Solved t' -> normalize t'
            _ -> typeVar <$> getKey' node
        Nothing -> pure t
    normalize t = pure t

    solveCt' :: (BreadCrumbs, (Type, Type)) -> SolveM s ()
    solveCt' (bcs, (t1, t2)) = do
      t1' <- normalize t1
      t2' <- normalize t2
      case (t1', t2') of
        ( Scalar (TypeVar _ (QualName [] v1) []),
          Scalar (TypeVar _ (QualName [] v2) [])
          )
            | v1 == v2 -> pure ()
            | otherwise -> do
                mb_node1 <- flexible v1
                mb_node2 <- flexible v2
                case (mb_node1, mb_node2) of
                  (Nothing, Nothing) ->
                    cannotUnify reason mempty bcs t1 t2
                  (Just v1_node, Nothing) ->
                    bindTyVar reason bcs v1 v1_node t2'
                  (Nothing, Just v2_node) ->
                    bindTyVar reason bcs v2 v2_node t1'
                  (Just v1_node, Just v2_node) ->
                    unionTyVars reason bcs v1 v1_node v2_node
        (Scalar (TypeVar _ (QualName [] v1) []), _) -> do
          mb_node <- flexible v1
          case mb_node of
            Just node -> bindTyVar reason bcs v1 node t2'
            Nothing -> tryUnify t1' t2' reason bcs
        (_, Scalar (TypeVar _ (QualName [] v2) [])) -> do
          mb_node <- flexible v2
          case mb_node of
            Just node -> bindTyVar reason bcs v2 node t1'
            Nothing -> tryUnify t1' t2' reason bcs
        (_, _) -> tryUnify t1' t2' reason bcs

    tryUnify :: Type -> Type -> Reason Type -> BreadCrumbs -> SolveM s ()
    tryUnify t1 t2 r bcs =
      case unify t1 t2 of
        Left details -> cannotUnify r (foldMap aNote details) bcs t1 t2
        Right eqs -> mapM_ solveCt' eqs

-- | Unify at the root, emitting new equalities that must hold.
unify :: Type -> Type -> Either (Maybe (Doc a)) [(BreadCrumbs, (Type, Type))]
unify (Scalar (Prim pt1)) (Scalar (Prim pt2))
  | pt1 == pt2 = Right []
unify
  (Scalar (TypeVar _ (QualName _ v1) targs1))
  (Scalar (TypeVar _ (QualName _ v2) targs2))
    | v1 == v2 =
        Right $ mapMaybe f $ zip targs1 targs2
    where
      f (TypeArgType t1, TypeArgType t2) = Just (mempty, (t1, t2))
      f _ = Nothing
unify
  (Scalar (Arrow _ _ _ t1a (RetType _ t1r)))
  (Scalar (Arrow _ _ _ t2a (RetType _ t2r))) =
    Right [(mempty, (t1a, t2a)), (mempty, (t1r', t2r'))]
    where
      t1r' = t1r `setUniqueness` NoUniqueness
      t2r' = t2r `setUniqueness` NoUniqueness
unify (Scalar (Record fs1)) (Scalar (Record fs2))
  | M.keys fs1 == M.keys fs2 =
      Right $
        map (first matchingField) $
          M.toList $
            M.intersectionWith (,) fs1 fs2
  | Just n1 <- length <$> areTupleFields fs1,
    Just n2 <- length <$> areTupleFields fs2,
    n1 /= n2 =
      Left . Just $
        "Tuples have"
          <+> pretty n1
          <+> "and"
          <+> pretty n2
          <+> "elements respectively."
  | otherwise =
      let missing =
            filter (`notElem` M.keys fs1) (M.keys fs2)
              <> filter (`notElem` M.keys fs2) (M.keys fs1)
       in Left . Just $
            "unshared fields:" <+> commasep (map pretty missing) <> "."
unify (Scalar (Sum cs1)) (Scalar (Sum cs2))
  | M.keys cs1 == M.keys cs2 =
      fmap concat . forM cs' $ \(c, (ts1, ts2)) -> do
        if length ts1 == length ts2
          then Right $ zipWith (curry (matchingConstructor c,)) ts1 ts2
          else Left Nothing
  | otherwise =
      Left . Just $ unsharedConstructorsMsg cs1 cs2
  where
    cs' = M.toList $ M.intersectionWith (,) cs1 cs2
unify t1 t2
  | Just t1' <- peelArray 1 t1,
    Just t2' <- peelArray 1 t2 =
      Right [(mempty, (t1', t2'))]
unify _ _ = Left Nothing

maybeLookupTyVarSol :: TyVar -> SolveM s (Maybe TyVarSol)
maybeLookupTyVarSol tv = do
  tyvars <- asks solverTyVars
  case M.lookup tv tyvars of
    Nothing -> pure Nothing
    Just node -> do
      sol <- getSol' node
      pure $ Just sol

lookupTyVar :: TyVar -> SolveM s (Either (TyVarInfo ()) Type)
lookupTyVar tv =
  maybe bad unpack <$> maybeLookupTyVarSol tv
  where
    bad = error $ "Unknown tyvar: " <> prettyNameString tv
    unpack (Param {}) = error $ "Is a type param: " <> prettyNameString tv
    unpack (Solved t) = Right t
    unpack (Unsolved info) = Left info

lookupTyVarInfo :: TyVarNode s -> SolveM s (TyVarInfo ())
lookupTyVarInfo v_node = do
  r <- getSol' v_node
  case r of
    Unsolved info -> pure info
    _ -> do
      v <- getKey' v_node
      error $ "Tyvar is nonflexible: " <> prettyNameString v

lookupUF :: TyVar -> SolveM s (TyVarNode s)
lookupUF tv = do
  uf <- asks solverTyVars
  case M.lookup tv uf of
    Nothing -> error $ "Unknown tyvar: " <> prettyNameString tv
    Just node -> pure node

unifySharedFields ::
  Reason Type ->
  BreadCrumbs ->
  M.Map Name Type ->
  M.Map Name Type ->
  SolveM s ()
unifySharedFields reason bcs fs1 fs2 =
  forM_ (M.toList $ M.intersectionWith (,) fs1 fs2) $ \(f, (ts1, ts2)) ->
    solveEq reason (matchingField f <> bcs) ts1 ts2

unifySharedConstructors ::
  Reason Type ->
  BreadCrumbs ->
  M.Map Name [Type] ->
  M.Map Name [Type] ->
  SolveM s ()
unifySharedConstructors reason bcs cs1 cs2 =
  forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (ts1, ts2)) ->
    if length ts1 == length ts2
      then zipWithM_ (solveEq reason $ matchingConstructor c <> bcs) ts1 ts2
      else
        typeError (locOf reason) mempty $
          "Cannot unify type with constructor"
            </> indent 2 (pretty (Sum (M.singleton c ts1)))
            </> "with type of constructor"
            </> indent 2 (pretty (Sum (M.singleton c ts2)))
            </> "because they differ in arity."

setInfo :: TyVarNode s -> TyVarSol -> SolveM s ()
setInfo node sol = liftST $ assignNewSol node sol

unionTyVars ::
  Reason Type ->
  BreadCrumbs ->
  VName ->
  TyVarNode s ->
  TyVarNode s ->
  SolveM s ()
unionTyVars reason bcs v v_node t_node = do
  v_sol <- getSol' v_node
  t_info <- lookupTyVarInfo t_node
  c <- check v_sol t_info
  case c of
    Left (loc, notes, msg) -> typeError loc notes msg
    Right (Just new_sol) -> unionNewSol' v_node t_node new_sol
    Right Nothing -> union' v_node t_node
  where
    check ::
      TyVarSol ->
      TyVarInfo () ->
      SolveM s (Either (Loc, Notes, Doc ()) (Maybe TyVarSol))
    check v_sol t_info =
      case (v_sol, t_info) of
        (Unsolved (TyVarFree _ v_l), TyVarFree t_loc t_l)
          | v_l /= t_l ->
              pure $ Right $ Just $ Unsolved $ TyVarFree t_loc (min v_l t_l)
        (Unsolved info, TyVarFree {}) -> do
          pure $ Right $ Just $ Unsolved info
        --
        -- TyVarPrim cases
        ( Unsolved (TyVarPrim _ v_pts),
          TyVarPrim t_loc t_pts
          ) ->
            let pts = L.intersect v_pts t_pts
             in case pts of
                  [] ->
                    pure $
                      Left
                        ( locOf reason,
                          mempty,
                          "Cannot unify type that must be one of"
                            </> indent 2 (pretty v_pts)
                            </> "with type that must be one of"
                            </> indent 2 (pretty t_pts)
                        )
                  _ -> pure $ Right $ Just $ Unsolved $ TyVarPrim t_loc pts
        (Unsolved (TyVarPrim _ v_pts), TyVarRecord {}) ->
          pure $
            Left
              ( locOf reason,
                mempty,
                "Cannot unify type that must be one of"
                  </> indent 2 (pretty v_pts)
                  </> "with type that must be a record."
              )
        (Unsolved (TyVarPrim _ v_pts), TyVarSum {}) ->
          pure $
            Left
              ( locOf reason,
                mempty,
                "Cannot unify type that must be one of"
                  </> indent 2 (pretty v_pts)
                  </> "with type that must be sum."
              )
        --
        -- TyVarSum cases
        ( Unsolved (TyVarSum _ cs1),
          TyVarSum loc cs2
          ) -> do
            unifySharedConstructors reason bcs cs1 cs2
            let cs3 = cs1 <> cs2
            pure $ Right $ Just $ Unsolved $ TyVarSum loc cs3
        ( Unsolved TyVarSum {},
          TyVarPrim _ pts
          ) ->
            pure $
              Left
                ( locOf reason,
                  mempty,
                  "A sum type cannot be one of"
                    </> indent 2 (pretty pts)
                )
        ( Unsolved (TyVarSum _ cs1),
          TyVarRecord _ fs
          ) ->
            pure $
              Left
                ( locOf reason,
                  mempty,
                  "Cannot unify type with constructors"
                    </> indent 2 (pretty (Sum cs1))
                    </> "with type"
                    </> indent 2 (pretty (Scalar (Record fs)))
                )
        --
        -- TyVarRecord cases
        ( Unsolved (TyVarRecord _ fs1),
          TyVarRecord loc fs2
          ) -> do
            unifySharedFields reason bcs fs1 fs2
            let fs3 = fs1 <> fs2
            pure $ Right $ Just $ Unsolved $ TyVarRecord loc fs3
        ( Unsolved TyVarRecord {},
          TyVarPrim _ pts
          ) ->
            pure $
              Left
                ( locOf reason,
                  mempty,
                  "A record type cannot be one of"
                    </> indent 2 (pretty pts)
                )
        ( Unsolved (TyVarRecord _ fs1),
          TyVarSum _ cs
          ) ->
            pure $
              Left
                ( locOf reason,
                  mempty,
                  "Cannot unify record type"
                    </> indent 2 (pretty (Record fs1))
                    </> "with type"
                    </> indent 2 (pretty (Scalar (Sum cs)))
                )
        --
        -- Internal error cases
        (Solved {}, _) -> alreadySolved
        (Param {}, _) -> isParam
        _ -> pure $ Right Nothing

    alreadySolved = error $ "Type variable already solved: " <> prettyNameString v
    isParam = error $ "Type name is a type parameter: " <> prettyNameString v

scopeViolation :: Reason Type -> VName -> Type -> VName -> SolveM s ()
scopeViolation reason v1 ty v2 =
  typeError (locOf reason) mempty . withIndexLink "scope-violation" $
    "Cannot unify type"
      </> indent 2 (pretty ty)
      </> "with"
      <+> dquotes (prettyName v1)
      <+> "(scope violation)."
      </> "This is because"
      <+> dquotes (prettyName v2)
      <+> "is rigidly bound in a deeper scope."

scopeCheck :: Reason Type -> TyVar -> Level -> Type -> SolveM s ()
scopeCheck reason v v_lvl ty = mapM_ check $ typeVars ty
  where
    check :: TyVar -> SolveM s ()
    check ty_v = do
      maybe (pure ()) checkNode =<< maybeLookupUF ty_v

    checkNode :: TyVarNode s -> SolveM s ()
    checkNode node = do
      sol <- getSol' node
      case sol of
        Param ty_v_lvl _ _
          | ty_v_lvl > v_lvl -> do
              k <- getKey' node
              ty' <- substTyVars ty
              scopeViolation reason v ty' k
        Solved ty' -> do
          mapM_ check $ typeVars ty'
        _ -> pure ()

-- | If a type variable has a liftedness constraint, we propagate that
-- constraint to its solution. The actual checking for correct usage
-- is done later, by 'localChecks' and 'instTyVars' in the sized type
-- checker, which know why the constraints exist and can produce
-- proper error messages.
liftednessCheck :: Liftedness -> Type -> SolveM s ()
liftednessCheck l (Scalar (TypeVar _ (QualName [] v) _)) = do
  v_info <- maybeLookupTyVarSol v
  case v_info of
    Nothing ->
      -- Is an opaque type.
      pure ()
    Just (Solved v_ty) ->
      liftednessCheck l v_ty
    Just Param {} -> pure ()
    Just (Unsolved (TyVarFree loc v_l))
      | l < v_l -> do
          node <- lookupUF v
          setInfo node $ Unsolved $ TyVarFree loc l
    Just Unsolved {} -> pure ()
liftednessCheck _ (Scalar Prim {}) = pure ()
liftednessCheck Lifted _ = pure ()
liftednessCheck _ Array {} = pure ()
liftednessCheck _ (Scalar Arrow {}) = pure ()
liftednessCheck l (Scalar (Record fs)) =
  mapM_ (liftednessCheck l) fs
liftednessCheck l (Scalar (Sum cs)) =
  mapM_ (mapM_ $ liftednessCheck l) cs
liftednessCheck _ (Scalar TypeVar {}) = pure ()

solveTyVar :: (VName, (Level, TyVarInfo ())) -> SolveM s ()
solveTyVar (tv, (lvl, TyVarFree loc l)) = do
  tv_t <- lookupTyVar tv
  case tv_t of
    Right ty -> do
      scopeCheck (Reason loc) tv lvl ty
      liftednessCheck l ty
    _ -> pure ()
solveTyVar (tv, (_, TyVarPrim loc pts)) = do
  tv_t <- lookupTyVar tv
  case tv_t of
    Right ty
      | ty `elem` map (Scalar . Prim) pts -> pure ()
      | otherwise ->
          typeError loc mempty $
            "Numeric constant inferred to be of type"
              </> indent 2 (align (pretty ty))
              </> "which is not possible."
    _ -> pure ()
solveTyVar (tv, (_, TyVarRecord loc fs1)) = do
  tv_t <- lookupTyVar tv
  case tv_t of
    Left _ ->
      typeError loc mempty . withIndexLink "ambiguous-type" $
        "Type"
          <+> prettyName tv
          <+> "is ambiguous."
          </> "Must be a record with fields"
          </> indent 2 (pretty (Scalar (Record fs1)))
    Right _ -> pure ()
solveTyVar (tv, (_, TyVarSum loc cs1)) = do
  tv_t <- lookupTyVar tv
  case tv_t of
    Left _ ->
      typeError loc mempty . withIndexLink "ambiguous-type" $
        "Type is ambiguous."
          </> "Must be a sum type with constructors"
          </> indent 2 (pretty (Scalar (Sum cs1)))
    Right _ -> pure ()

maybeLookupUF :: TyVar -> SolveM s (Maybe (TyVarNode s))
maybeLookupUF tv = do
  uf <- asks solverTyVars
  pure . M.lookup tv $ uf

getSolution :: SolveM s ([UnconTyVar], Solution)
getSolution = do
  uf <- asks solverTyVars
  resolved <- M.traverseWithKey resolve uf
  let unconstrained = M.foldrWithKey unconstr [] resolved
      sol = M.mapMaybeWithKey mkSubst resolved
  pure (unconstrained, sol)
  where
    resolve ::
      TyVar ->
      TyVarNode s ->
      SolveM s (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness)
    resolve tv node = do
      sol <- getSol' node
      case sol of
        Unsolved (TyVarFree _ l) -> do
          k <- getKey' node
          let tv' = typeVar k
          -- If the current type variable and root type variable are
          -- different, this variable is unconstrained, so we save the
          -- liftedness constraint for later.
          pure (Right tv', if k == tv then Just l else Nothing)
        Unsolved (TyVarPrim _ pts) -> pure (Left pts, Nothing)
        Solved t -> do
          t' <- substTyVars t
          pure (Right $ first (const ()) t', Nothing)
        _ -> do
          k <- getKey' node
          pure (Right $ typeVar k, Nothing)

    unconstr ::
      TyVar ->
      (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness) ->
      [UnconTyVar] ->
      [UnconTyVar]
    unconstr tv (_, Just l) acc = (tv, l) : acc
    unconstr _ _ acc = acc

    mkSubst ::
      TyVar ->
      (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness) ->
      Maybe (Either [PrimType] (TypeBase () NoUniqueness))
    mkSubst _ (_, Just _) = Nothing
    mkSubst tv (s@(Right (Scalar (TypeVar _ (QualName [] tv') _))), _) =
      if tv /= tv' then Just s else Nothing
    mkSubst _ (s, _) = Just s

-- | Solve type constraints, producing either an error or a solution,
-- alongside a list of unconstrained type variables.
solve ::
  [CtTy ()] ->
  TyParams ->
  TyVars () ->
  Either TypeError ([UnconTyVar], Solution)
solve constraints typarams tyvars =
  runST $ do
    r <- initializeState typarams tyvars
    flip runReaderT r $ runExceptT $ runSolveM $ do
      mapM_ solveCt constraints
      mapM_ solveTyVar $ M.toList tyvars
      getSolution
{-# NOINLINE solve #-}