-- | Constraints produced (and solved) by the type checker.
module Language.Futhark.TypeChecker.Constraints
  ( Reason (..),
    CtType,
    CtTy (..),
    TyVarInfo (..),
    Level,
    TyVar,
    TyVars,
    TyParams,
  )
where

import Data.Bifunctor
import Data.Loc
import Data.Map qualified as M
import Futhark.Util.Pretty
import Language.Futhark

-- | The type representation used by the constraint solver.
type CtType d = TypeBase d NoMode

-- | The reason for a type constraint. Used to generate type error
-- messages. The expected type is always the first one.
data Reason t
  = -- | No particular reason.
    Reason Loc
  | -- | Arising from pattern match.
    ReasonPatMatch Loc (PatBase NoInfo VName StructType) t
  | -- | Arising from explicit ascription.
    ReasonAscription Loc t t
  | ReasonRetType Loc t t
  | -- | Arising from checking a function argument. The types are the
    -- expected (parameter) and actual (argument) types.
    ReasonApply Loc (Maybe (QualName VName), Int) Exp t t
  | -- | Used when unifying a type with a function type in a function
    -- application. If this unification fails, it means the supposed
    -- function was not a function after all.
    ReasonApplySplit Loc (Maybe (QualName VName), Int) Exp t
  | ReasonBranches Loc t t
  deriving (Reason t -> Reason t -> Bool
(Reason t -> Reason t -> Bool)
-> (Reason t -> Reason t -> Bool) -> Eq (Reason t)
forall t. Eq t => Reason t -> Reason t -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall t. Eq t => Reason t -> Reason t -> Bool
== :: Reason t -> Reason t -> Bool
$c/= :: forall t. Eq t => Reason t -> Reason t -> Bool
/= :: Reason t -> Reason t -> Bool
Eq, Int -> Reason t -> ShowS
[Reason t] -> ShowS
Reason t -> String
(Int -> Reason t -> ShowS)
-> (Reason t -> String) -> ([Reason t] -> ShowS) -> Show (Reason t)
forall t. Show t => Int -> Reason t -> ShowS
forall t. Show t => [Reason t] -> ShowS
forall t. Show t => Reason t -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall t. Show t => Int -> Reason t -> ShowS
showsPrec :: Int -> Reason t -> ShowS
$cshow :: forall t. Show t => Reason t -> String
show :: Reason t -> String
$cshowList :: forall t. Show t => [Reason t] -> ShowS
showList :: [Reason t] -> ShowS
Show, (forall a b. (a -> b) -> Reason a -> Reason b)
-> (forall a b. a -> Reason b -> Reason a) -> Functor Reason
forall a b. a -> Reason b -> Reason a
forall a b. (a -> b) -> Reason a -> Reason b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Reason a -> Reason b
fmap :: forall a b. (a -> b) -> Reason a -> Reason b
$c<$ :: forall a b. a -> Reason b -> Reason a
<$ :: forall a b. a -> Reason b -> Reason a
Functor, (forall m. Monoid m => Reason m -> m)
-> (forall m a. Monoid m => (a -> m) -> Reason a -> m)
-> (forall m a. Monoid m => (a -> m) -> Reason a -> m)
-> (forall a b. (a -> b -> b) -> b -> Reason a -> b)
-> (forall a b. (a -> b -> b) -> b -> Reason a -> b)
-> (forall b a. (b -> a -> b) -> b -> Reason a -> b)
-> (forall b a. (b -> a -> b) -> b -> Reason a -> b)
-> (forall a. (a -> a -> a) -> Reason a -> a)
-> (forall a. (a -> a -> a) -> Reason a -> a)
-> (forall a. Reason a -> [a])
-> (forall a. Reason a -> Bool)
-> (forall a. Reason a -> Int)
-> (forall a. Eq a => a -> Reason a -> Bool)
-> (forall a. Ord a => Reason a -> a)
-> (forall a. Ord a => Reason a -> a)
-> (forall a. Num a => Reason a -> a)
-> (forall a. Num a => Reason a -> a)
-> Foldable Reason
forall a. Eq a => a -> Reason a -> Bool
forall a. Num a => Reason a -> a
forall a. Ord a => Reason a -> a
forall m. Monoid m => Reason m -> m
forall a. Reason a -> Bool
forall a. Reason a -> Int
forall a. Reason a -> [a]
forall a. (a -> a -> a) -> Reason a -> a
forall m a. Monoid m => (a -> m) -> Reason a -> m
forall b a. (b -> a -> b) -> b -> Reason a -> b
forall a b. (a -> b -> b) -> b -> Reason a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Reason m -> m
fold :: forall m. Monoid m => Reason m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Reason a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Reason a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Reason a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Reason a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Reason a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Reason a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Reason a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Reason a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Reason a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Reason a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Reason a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Reason a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Reason a -> a
foldr1 :: forall a. (a -> a -> a) -> Reason a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Reason a -> a
foldl1 :: forall a. (a -> a -> a) -> Reason a -> a
$ctoList :: forall a. Reason a -> [a]
toList :: forall a. Reason a -> [a]
$cnull :: forall a. Reason a -> Bool
null :: forall a. Reason a -> Bool
$clength :: forall a. Reason a -> Int
length :: forall a. Reason a -> Int
$celem :: forall a. Eq a => a -> Reason a -> Bool
elem :: forall a. Eq a => a -> Reason a -> Bool
$cmaximum :: forall a. Ord a => Reason a -> a
maximum :: forall a. Ord a => Reason a -> a
$cminimum :: forall a. Ord a => Reason a -> a
minimum :: forall a. Ord a => Reason a -> a
$csum :: forall a. Num a => Reason a -> a
sum :: forall a. Num a => Reason a -> a
$cproduct :: forall a. Num a => Reason a -> a
product :: forall a. Num a => Reason a -> a
Foldable, Functor Reason
Foldable Reason
(Functor Reason, Foldable Reason) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Reason a -> f (Reason b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Reason (f a) -> f (Reason a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Reason a -> m (Reason b))
-> (forall (m :: * -> *) a.
    Monad m =>
    Reason (m a) -> m (Reason a))
-> Traversable Reason
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Reason (m a) -> m (Reason a)
forall (f :: * -> *) a.
Applicative f =>
Reason (f a) -> f (Reason a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Reason a -> m (Reason b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Reason a -> f (Reason b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Reason a -> f (Reason b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Reason a -> f (Reason b)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
Reason (f a) -> f (Reason a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Reason (f a) -> f (Reason a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Reason a -> m (Reason b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Reason a -> m (Reason b)
$csequence :: forall (m :: * -> *) a. Monad m => Reason (m a) -> m (Reason a)
sequence :: forall (m :: * -> *) a. Monad m => Reason (m a) -> m (Reason a)
Traversable)

instance Located (Reason t) where
  locOf :: Reason t -> Loc
locOf (Reason Loc
l) = Loc
l
  locOf (ReasonPatMatch Loc
l PatBase NoInfo VName StructType
_ t
_) = Loc
l
  locOf (ReasonAscription Loc
l t
_ t
_) = Loc
l
  locOf (ReasonRetType Loc
l t
_ t
_) = Loc
l
  locOf (ReasonApply Loc
l (Maybe (QualName VName), Int)
_ Exp
_ t
_ t
_) = Loc
l
  locOf (ReasonApplySplit Loc
l (Maybe (QualName VName), Int)
_ Exp
_ t
_) = Loc
l
  locOf (ReasonBranches Loc
l t
_ t
_) = Loc
l

-- | A type constraint.
data CtTy d = CtEq (Reason (CtType d)) (TypeBase d NoMode) (TypeBase d NoMode)
  deriving (Int -> CtTy d -> ShowS
[CtTy d] -> ShowS
CtTy d -> String
(Int -> CtTy d -> ShowS)
-> (CtTy d -> String) -> ([CtTy d] -> ShowS) -> Show (CtTy d)
forall d. Show d => Int -> CtTy d -> ShowS
forall d. Show d => [CtTy d] -> ShowS
forall d. Show d => CtTy d -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall d. Show d => Int -> CtTy d -> ShowS
showsPrec :: Int -> CtTy d -> ShowS
$cshow :: forall d. Show d => CtTy d -> String
show :: CtTy d -> String
$cshowList :: forall d. Show d => [CtTy d] -> ShowS
showList :: [CtTy d] -> ShowS
Show)

instance Functor CtTy where
  fmap :: forall a b. (a -> b) -> CtTy a -> CtTy b
fmap a -> b
f (CtEq Reason (CtType a)
r CtType a
x CtType a
y) = Reason (CtType b) -> CtType b -> CtType b -> CtTy b
forall d. Reason (CtType d) -> CtType d -> CtType d -> CtTy d
CtEq ((CtType a -> CtType b) -> Reason (CtType a) -> Reason (CtType b)
forall a b. (a -> b) -> Reason a -> Reason b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> CtType a -> CtType b
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f) Reason (CtType a)
r) ((a -> b) -> CtType a -> CtType b
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f CtType a
x) ((a -> b) -> CtType a -> CtType b
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f CtType a
y)

ctReason :: CtTy d -> Reason (CtType d)
ctReason :: forall d. CtTy d -> Reason (CtType d)
ctReason (CtEq Reason (CtType d)
r CtType d
_ CtType d
_) = Reason (CtType d)
r

instance Located (CtTy d) where
  locOf :: CtTy d -> Loc
locOf = Reason (CtType d) -> Loc
forall a. Located a => a -> Loc
locOf (Reason (CtType d) -> Loc)
-> (CtTy d -> Reason (CtType d)) -> CtTy d -> Loc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CtTy d -> Reason (CtType d)
forall d. CtTy d -> Reason (CtType d)
ctReason

instance Pretty (CtTy Size) where
  pretty :: forall ann. CtTy Exp -> Doc ann
pretty (CtEq Reason StructType
_ StructType
t1 StructType
t2) = StructType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
t1 Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"~" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> StructType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
t2

instance Pretty (CtTy ()) where
  pretty :: forall ann. CtTy () -> Doc ann
pretty (CtEq Reason (TypeBase () NoMode)
_ TypeBase () NoMode
t1 TypeBase () NoMode
t2) = TypeBase () NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase () NoMode -> Doc ann
pretty TypeBase () NoMode
t1 Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"~" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> TypeBase () NoMode -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase () NoMode -> Doc ann
pretty TypeBase () NoMode
t2

-- | Information about a flexible type variable. Every type variable
-- is associated with a location, which is the original syntax element
-- that it is the type of.
data TyVarInfo d
  = -- | Can be substituted with anything.
    TyVarFree Loc Liftedness
  | -- | Can only be substituted with these primitive types.
    TyVarPrim Loc [PrimType]
  | -- | Must be a record with these fields.
    TyVarRecord Loc (M.Map Name (CtType d))
  | -- | Must be a sum type with these fields.
    TyVarSum Loc (M.Map Name [CtType d])
  deriving (Int -> TyVarInfo d -> ShowS
[TyVarInfo d] -> ShowS
TyVarInfo d -> String
(Int -> TyVarInfo d -> ShowS)
-> (TyVarInfo d -> String)
-> ([TyVarInfo d] -> ShowS)
-> Show (TyVarInfo d)
forall d. Show d => Int -> TyVarInfo d -> ShowS
forall d. Show d => [TyVarInfo d] -> ShowS
forall d. Show d => TyVarInfo d -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall d. Show d => Int -> TyVarInfo d -> ShowS
showsPrec :: Int -> TyVarInfo d -> ShowS
$cshow :: forall d. Show d => TyVarInfo d -> String
show :: TyVarInfo d -> String
$cshowList :: forall d. Show d => [TyVarInfo d] -> ShowS
showList :: [TyVarInfo d] -> ShowS
Show, TyVarInfo d -> TyVarInfo d -> Bool
(TyVarInfo d -> TyVarInfo d -> Bool)
-> (TyVarInfo d -> TyVarInfo d -> Bool) -> Eq (TyVarInfo d)
forall d. Eq d => TyVarInfo d -> TyVarInfo d -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall d. Eq d => TyVarInfo d -> TyVarInfo d -> Bool
== :: TyVarInfo d -> TyVarInfo d -> Bool
$c/= :: forall d. Eq d => TyVarInfo d -> TyVarInfo d -> Bool
/= :: TyVarInfo d -> TyVarInfo d -> Bool
Eq)

instance Functor TyVarInfo where
  fmap :: forall a b. (a -> b) -> TyVarInfo a -> TyVarInfo b
fmap a -> b
_ (TyVarFree Loc
loc Liftedness
l) = Loc -> Liftedness -> TyVarInfo b
forall d. Loc -> Liftedness -> TyVarInfo d
TyVarFree Loc
loc Liftedness
l
  fmap a -> b
_ (TyVarPrim Loc
loc [PrimType]
ts) = Loc -> [PrimType] -> TyVarInfo b
forall d. Loc -> [PrimType] -> TyVarInfo d
TyVarPrim Loc
loc [PrimType]
ts
  fmap a -> b
f (TyVarRecord Loc
loc Map Name (CtType a)
m) = Loc -> Map Name (CtType b) -> TyVarInfo b
forall d. Loc -> Map Name (CtType d) -> TyVarInfo d
TyVarRecord Loc
loc (Map Name (CtType b) -> TyVarInfo b)
-> Map Name (CtType b) -> TyVarInfo b
forall a b. (a -> b) -> a -> b
$ (CtType a -> CtType b)
-> Map Name (CtType a) -> Map Name (CtType b)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((a -> b) -> CtType a -> CtType b
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f) Map Name (CtType a)
m
  fmap a -> b
f (TyVarSum Loc
loc Map Name [CtType a]
m) = Loc -> Map Name [CtType b] -> TyVarInfo b
forall d. Loc -> Map Name [CtType d] -> TyVarInfo d
TyVarSum Loc
loc (Map Name [CtType b] -> TyVarInfo b)
-> Map Name [CtType b] -> TyVarInfo b
forall a b. (a -> b) -> a -> b
$ ([CtType a] -> [CtType b])
-> Map Name [CtType a] -> Map Name [CtType b]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((CtType a -> CtType b) -> [CtType a] -> [CtType b]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> CtType a -> CtType b
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f)) Map Name [CtType a]
m

prettyTyVarInfo :: (Pretty (Shape d)) => TyVarInfo d -> Doc a
prettyTyVarInfo :: forall d a. Pretty (Shape d) => TyVarInfo d -> Doc a
prettyTyVarInfo (TyVarFree Loc
_ Liftedness
l) = Doc a
"free" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Liftedness -> Doc a
forall a ann. Pretty a => a -> Doc ann
forall ann. Liftedness -> Doc ann
pretty Liftedness
l
prettyTyVarInfo (TyVarPrim Loc
_ [PrimType]
pts) = Doc a
"∈" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [PrimType] -> Doc a
forall ann. [PrimType] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [PrimType]
pts
prettyTyVarInfo (TyVarRecord Loc
_ Map Name (TypeBase d NoMode)
fs) = TypeBase d NoMode -> Doc a
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase d NoMode -> Doc ann
pretty (TypeBase d NoMode -> Doc a) -> TypeBase d NoMode -> Doc a
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase d NoMode -> TypeBase d NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase d NoMode -> TypeBase d NoMode)
-> ScalarTypeBase d NoMode -> TypeBase d NoMode
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase d NoMode) -> ScalarTypeBase d NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record Map Name (TypeBase d NoMode)
fs
prettyTyVarInfo (TyVarSum Loc
_ Map Name [TypeBase d NoMode]
cs) = TypeBase d NoMode -> Doc a
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase d NoMode -> Doc ann
pretty (TypeBase d NoMode -> Doc a) -> TypeBase d NoMode -> Doc a
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase d NoMode -> TypeBase d NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase d NoMode -> TypeBase d NoMode)
-> ScalarTypeBase d NoMode -> TypeBase d NoMode
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase d NoMode] -> ScalarTypeBase d NoMode
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum Map Name [TypeBase d NoMode]
cs

instance Pretty (TyVarInfo ()) where
  pretty :: forall ann. TyVarInfo () -> Doc ann
pretty = TyVarInfo () -> Doc ann
forall d a. Pretty (Shape d) => TyVarInfo d -> Doc a
prettyTyVarInfo

instance Located (TyVarInfo d) where
  locOf :: TyVarInfo d -> Loc
locOf (TyVarFree Loc
loc Liftedness
_) = Loc
loc
  locOf (TyVarPrim Loc
loc [PrimType]
_) = Loc
loc
  locOf (TyVarRecord Loc
loc Map Name (CtType d)
_) = Loc
loc
  locOf (TyVarSum Loc
loc Map Name [CtType d]
_) = Loc
loc

-- | The name of a type variable.
type TyVar = VName

-- | The level at which a type variable is bound.  Higher means
-- deeper.  We can only unify a type variable at level @i@ with a type
-- @t@ if all type names that occur in @t@ are at most at level @i@.
type Level = Int

-- | If a VName is not in this map, it should be in the 'TyParams' -
-- the exception is abstract types, which are just missing (and
-- assumed to have smallest possible level).
type TyVars d = M.Map TyVar (Level, TyVarInfo d)

-- | Explicit type parameters.
type TyParams = M.Map TyVar (Level, Liftedness, Loc)