Types.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}

-- | Type checker building blocks that do not involve unification.
module Language.Futhark.TypeChecker.Types
  ( checkTypeExp,
    checkTypeDecl,
    unifyTypesU,
    subtypeOf,
    subuniqueOf,
    checkForDuplicateNames,
    checkTypeParams,
    typeParamToArg,
    TypeSub (..),
    TypeSubs,
    substituteTypes,
    Subst (..),
    Substitutable (..),
    substTypesAny,
  )
where

import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.List (foldl', nub, sort)
import qualified Data.Map.Strict as M
import Data.Maybe
import Futhark.Util.Pretty
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Monad

-- | @unifyTypes uf t1 t2@ attempts to unify @t1@ and @t2@.  If
-- unification cannot happen, 'Nothing' is returned, otherwise a type
-- that combines the aliasing of @t1@ and @t2@ is returned.
-- Uniqueness is unified with @uf@.
unifyTypesU ::
  (Monoid als, ArrayDim dim) =>
  (Uniqueness -> Uniqueness -> Maybe Uniqueness) ->
  TypeBase dim als ->
  TypeBase dim als ->
  Maybe (TypeBase dim als)
unifyTypesU uf (Array als1 u1 et1 shape1) (Array als2 u2 et2 shape2) =
  Array (als1 <> als2) <$> uf u1 u2
    <*> unifyScalarTypes uf et1 et2
    <*> unifyShapes shape1 shape2
unifyTypesU uf (Scalar t1) (Scalar t2) = Scalar <$> unifyScalarTypes uf t1 t2
unifyTypesU _ _ _ = Nothing

unifyScalarTypes ::
  (Monoid als, ArrayDim dim) =>
  (Uniqueness -> Uniqueness -> Maybe Uniqueness) ->
  ScalarTypeBase dim als ->
  ScalarTypeBase dim als ->
  Maybe (ScalarTypeBase dim als)
unifyScalarTypes _ (Prim t1) (Prim t2)
  | t1 == t2 = Just $ Prim t1
  | otherwise = Nothing
unifyScalarTypes uf (TypeVar als1 u1 t1 targs1) (TypeVar als2 u2 t2 targs2)
  | t1 == t2 = do
    u3 <- uf u1 u2
    targs3 <- zipWithM (unifyTypeArgs uf) targs1 targs2
    Just $ TypeVar (als1 <> als2) u3 t1 targs3
  | otherwise = Nothing
unifyScalarTypes uf (Record ts1) (Record ts2)
  | length ts1 == length ts2,
    sort (M.keys ts1) == sort (M.keys ts2) =
    Record
      <$> traverse
        (uncurry (unifyTypesU uf))
        (M.intersectionWith (,) ts1 ts2)
unifyScalarTypes uf (Arrow as1 mn1 t1 t1') (Arrow as2 _ t2 t2') =
  Arrow (as1 <> as2) mn1 <$> unifyTypesU (flip uf) t1 t2 <*> unifyTypesU uf t1' t2'
unifyScalarTypes uf (Sum cs1) (Sum cs2)
  | length cs1 == length cs2,
    sort (M.keys cs1) == sort (M.keys cs2) =
    Sum
      <$> traverse
        (uncurry (zipWithM (unifyTypesU uf)))
        (M.intersectionWith (,) cs1 cs2)
unifyScalarTypes _ _ _ = Nothing

unifyTypeArgs ::
  (ArrayDim dim) =>
  (Uniqueness -> Uniqueness -> Maybe Uniqueness) ->
  TypeArg dim ->
  TypeArg dim ->
  Maybe (TypeArg dim)
unifyTypeArgs _ (TypeArgDim d1 loc) (TypeArgDim d2 _) =
  TypeArgDim <$> unifyDims d1 d2 <*> pure loc
unifyTypeArgs uf (TypeArgType t1 loc) (TypeArgType t2 _) =
  TypeArgType <$> unifyTypesU uf t1 t2 <*> pure loc
unifyTypeArgs _ _ _ =
  Nothing

-- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to
-- @y@), meaning @x@ is valid whenever @y@ is.
subtypeOf ::
  ArrayDim dim =>
  TypeBase dim as1 ->
  TypeBase dim as2 ->
  Bool
subtypeOf t1 t2 = isJust $ unifyTypesU unifyUniqueness (toStruct t1) (toStruct t2)
  where
    unifyUniqueness u2 u1 = if u2 `subuniqueOf` u1 then Just u1 else Nothing

-- | @x `subuniqueOf` y@ is true if @x@ is not less unique than @y@.
subuniqueOf :: Uniqueness -> Uniqueness -> Bool
subuniqueOf Nonunique Unique = False
subuniqueOf _ _ = True

-- | Use 'checkTypeExp' to check a type declaration.
checkTypeDecl ::
  MonadTypeChecker m =>
  TypeDeclBase NoInfo Name ->
  m (TypeDeclBase Info VName, Liftedness)
checkTypeDecl (TypeDecl t NoInfo) = do
  checkForDuplicateNamesInType t
  (t', st, l) <- checkTypeExp t
  return (TypeDecl t' $ Info st, l)

-- | Type-check a single 'TypeExp', returning the checked 'TypeExp',
-- its fully expanded type (modulo yet-unelaborated type variables),
-- and whether it is potentially higher-order.
checkTypeExp ::
  MonadTypeChecker m =>
  TypeExp Name ->
  m (TypeExp VName, StructType, Liftedness)
checkTypeExp (TEVar name loc) = do
  (name', ps, t, l) <- lookupType loc name
  case ps of
    [] -> return (TEVar name' loc, t, l)
    _ ->
      typeError loc mempty $
        "Type constructor" <+> pquote (spread (ppr name : map ppr ps))
          <+> "used without any arguments."
checkTypeExp (TETuple ts loc) = do
  (ts', ts_s, ls) <- unzip3 <$> mapM checkTypeExp ts
  return (TETuple ts' loc, tupleRecord ts_s, foldl' max Unlifted ls)
checkTypeExp t@(TERecord fs loc) = do
  -- Check for duplicate field names.
  let field_names = map fst fs
  unless (sort field_names == sort (nub field_names)) $
    typeError loc mempty $ "Duplicate record fields in" <+> ppr t <> "."

  fs_ts_ls <- traverse checkTypeExp $ M.fromList fs
  let fs' = fmap (\(x, _, _) -> x) fs_ts_ls
      ts_s = fmap (\(_, y, _) -> y) fs_ts_ls
      ls = fmap (\(_, _, z) -> z) fs_ts_ls
  return
    ( TERecord (M.toList fs') loc,
      Scalar $ Record ts_s,
      foldl' max Unlifted ls
    )
checkTypeExp (TEArray t d loc) = do
  (t', st, l) <- checkTypeExp t
  (d', d'') <- checkDimExp d
  case (l, arrayOf st (ShapeDecl [d'']) Nonunique) of
    (Unlifted, st') -> return (TEArray t' d' loc, st', Unlifted)
    (SizeLifted, _) ->
      typeError loc mempty $
        "Cannot create array with elements of size-lifted type" <+> pquote (ppr t)
          <+/> "(might cause irregular array)."
    (Lifted, _) ->
      typeError loc mempty $
        "Cannot create array with elements of lifted type" <+> pquote (ppr t)
          <+/> "(might contain function)."
  where
    checkDimExp DimExpAny =
      return (DimExpAny, AnyDim)
    checkDimExp (DimExpConst k dloc) =
      return (DimExpConst k dloc, ConstDim k)
    checkDimExp (DimExpNamed v dloc) = do
      v' <- checkNamedDim loc v
      return (DimExpNamed v' dloc, NamedDim v')
checkTypeExp (TEUnique t loc) = do
  (t', st, l) <- checkTypeExp t
  unless (mayContainArray st) $
    warn loc $ "Declaring" <+> pquote (ppr st) <+> "as unique has no effect."
  return (TEUnique t' loc, st `setUniqueness` Unique, l)
  where
    mayContainArray (Scalar Prim {}) = False
    mayContainArray Array {} = True
    mayContainArray (Scalar (Record fs)) = any mayContainArray fs
    mayContainArray (Scalar TypeVar {}) = True
    mayContainArray (Scalar Arrow {}) = False
    mayContainArray (Scalar (Sum cs)) = (any . any) mayContainArray cs
checkTypeExp (TEArrow (Just v) t1 t2 loc) = do
  (t1', st1, _) <- checkTypeExp t1
  bindSpaced [(Term, v)] $ do
    v' <- checkName Term v loc
    bindVal v' (BoundV [] st1) $ do
      (t2', st2, _) <- checkTypeExp t2
      return
        ( TEArrow (Just v') t1' t2' loc,
          Scalar $ Arrow mempty (Named v') st1 st2,
          Lifted
        )
checkTypeExp (TEArrow Nothing t1 t2 loc) = do
  (t1', st1, _) <- checkTypeExp t1
  (t2', st2, _) <- checkTypeExp t2
  return
    ( TEArrow Nothing t1' t2' loc,
      Scalar $ Arrow mempty Unnamed st1 st2,
      Lifted
    )
checkTypeExp ote@TEApply {} = do
  (tname, tname_loc, targs) <- rootAndArgs ote
  (tname', ps, t, l) <- lookupType tloc tname
  if length ps /= length targs
    then
      typeError tloc mempty $
        "Type constructor" <+> pquote (ppr tname) <+> "requires" <+> ppr (length ps)
          <+> "arguments, but provided"
          <+> ppr (length targs) <> "."
    else do
      (targs', substs) <- unzip <$> zipWithM checkArgApply ps targs
      return
        ( foldl (\x y -> TEApply x y tloc) (TEVar tname' tname_loc) targs',
          substituteTypes (mconcat substs) t,
          l
        )
  where
    tloc = srclocOf ote

    rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
    rootAndArgs (TEVar qn loc) = return (qn, loc, [])
    rootAndArgs (TEApply op arg _) = do
      (op', loc, args) <- rootAndArgs op
      return (op', loc, args ++ [arg])
    rootAndArgs te' =
      typeError (srclocOf te') mempty $
        "Type" <+> pquote (ppr te') <+> "is not a type constructor."

    checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpNamed v dloc) loc) = do
      v' <- checkNamedDim loc v
      return
        ( TypeArgExpDim (DimExpNamed v' dloc) loc,
          M.singleton pv $ DimSub $ NamedDim v'
        )
    checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpConst x dloc) loc) =
      return
        ( TypeArgExpDim (DimExpConst x dloc) loc,
          M.singleton pv $ DimSub $ ConstDim x
        )
    checkArgApply (TypeParamDim pv _) (TypeArgExpDim DimExpAny loc) =
      return
        ( TypeArgExpDim DimExpAny loc,
          M.singleton pv $ DimSub AnyDim
        )
    checkArgApply (TypeParamType l pv _) (TypeArgExpType te) = do
      (te', st, _) <- checkTypeExp te
      return
        ( TypeArgExpType te',
          M.singleton pv $ TypeSub $ TypeAbbr l [] st
        )
    checkArgApply p a =
      typeError tloc mempty $
        "Type argument" <+> ppr a
          <+> "not valid for a type parameter"
          <+> ppr p <> "."
checkTypeExp t@(TESum cs loc) = do
  let constructors = map fst cs
  unless (sort constructors == sort (nub constructors)) $
    typeError loc mempty $ "Duplicate constructors in" <+> ppr t

  unless (length constructors < 256) $
    typeError loc mempty "Sum types must have less than 256 constructors."

  cs_ts_ls <- (traverse . traverse) checkTypeExp $ M.fromList cs
  let cs' = (fmap . fmap) (\(x, _, _) -> x) cs_ts_ls
      ts_s = (fmap . fmap) (\(_, y, _) -> y) cs_ts_ls
      ls = (concatMap . fmap) (\(_, _, z) -> z) cs_ts_ls
  return
    ( TESum (M.toList cs') loc,
      Scalar $ Sum ts_s,
      foldl' max Unlifted ls
    )

-- | Check for duplication of names inside a pattern group.  Produces
-- a description of all names used in the pattern group.
checkForDuplicateNames ::
  MonadTypeChecker m =>
  [UncheckedPattern] ->
  m ()
checkForDuplicateNames = (`evalStateT` mempty) . mapM_ check
  where
    check (Id v _ loc) = seen v loc
    check (PatternParens p _) = check p
    check Wildcard {} = return ()
    check (TuplePattern ps _) = mapM_ check ps
    check (RecordPattern fs _) = mapM_ (check . snd) fs
    check (PatternAscription p _ _) = check p
    check PatternLit {} = return ()
    check (PatternConstr _ _ ps _) = mapM_ check ps

    seen v loc = do
      already <- gets $ M.lookup v
      case already of
        Just prev_loc ->
          lift $
            typeError loc mempty $
              "Name" <+> pquote (ppr v) <+> "also bound at"
                <+> text (locStr prev_loc) <> "."
        Nothing ->
          modify $ M.insert v loc

-- | Check whether the type contains arrow types that define the same
-- parameter.  These might also exist further down, but that's not
-- really a problem - we mostly do this checking to help the user,
-- since it is likely an error, but it's easy to assign a semantics to
-- it (normal name shadowing).
checkForDuplicateNamesInType ::
  MonadTypeChecker m =>
  TypeExp Name ->
  m ()
checkForDuplicateNamesInType = check mempty
  where
    check seen (TEArrow (Just v) t1 t2 loc)
      | Just prev_loc <- M.lookup v seen =
        typeError loc mempty $
          text "Name" <+> pquote (ppr v)
            <+> "also bound at"
            <+> text (locStr prev_loc) <> "."
      | otherwise =
        check seen' t1 >> check seen' t2
      where
        seen' = M.insert v loc seen
    check seen (TEArrow Nothing t1 t2 _) =
      check seen t1 >> check seen t2
    check seen (TETuple ts _) = mapM_ (check seen) ts
    check seen (TERecord fs _) = mapM_ (check seen . snd) fs
    check seen (TEUnique t _) = check seen t
    check seen (TESum cs _) = mapM_ (mapM (check seen) . snd) cs
    check seen (TEApply t1 (TypeArgExpType t2) _) =
      check seen t1 >> check seen t2
    check seen (TEApply t1 TypeArgExpDim {} _) =
      check seen t1
    check _ TEArray {} = return ()
    check _ TEVar {} = return ()

-- | @checkTypeParams ps m@ checks the type parameters @ps@, then
-- invokes the continuation @m@ with the checked parameters, while
-- extending the monadic name map with @ps@.
checkTypeParams ::
  MonadTypeChecker m =>
  [TypeParamBase Name] ->
  ([TypeParamBase VName] -> m a) ->
  m a
checkTypeParams ps m =
  bindSpaced (map typeParamSpace ps) $
    m =<< evalStateT (mapM checkTypeParam ps) mempty
  where
    typeParamSpace (TypeParamDim pv _) = (Term, pv)
    typeParamSpace (TypeParamType _ pv _) = (Type, pv)

    checkParamName ns v loc = do
      seen <- gets $ M.lookup (ns, v)
      case seen of
        Just prev ->
          lift $
            typeError loc mempty $
              text "Type parameter" <+> pquote (ppr v)
                <+> "previously defined at"
                <+> text (locStr prev) <> "."
        Nothing -> do
          modify $ M.insert (ns, v) loc
          lift $ checkName ns v loc

    checkTypeParam (TypeParamDim pv loc) =
      TypeParamDim <$> checkParamName Term pv loc <*> pure loc
    checkTypeParam (TypeParamType l pv loc) =
      TypeParamType l <$> checkParamName Type pv loc <*> pure loc

-- | Construct a type argument corresponding to a type parameter.
typeParamToArg :: TypeParam -> StructTypeArg
typeParamToArg (TypeParamDim v ploc) =
  TypeArgDim (NamedDim $ qualName v) ploc
typeParamToArg (TypeParamType _ v ploc) =
  TypeArgType (Scalar $ TypeVar () Nonunique (typeName v) []) ploc

-- | A substitution for when using 'substituteTypes'.
data TypeSub
  = TypeSub TypeBinding
  | DimSub (DimDecl VName)
  deriving (Show)

-- | A collection of type substitutions.
type TypeSubs = M.Map VName TypeSub

-- | Apply type substitutions to the given type.
substituteTypes :: Monoid als => TypeSubs -> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes substs ot = case ot of
  Array als u at shape ->
    arrayOf
      (substituteTypes substs (Scalar at) `setAliases` mempty)
      (substituteInShape shape)
      u
      `addAliases` (<> als)
  Scalar (Prim t) -> Scalar $ Prim t
  Scalar (TypeVar als u v targs)
    | Just (TypeSub (TypeAbbr _ ps t)) <-
        M.lookup (qualLeaf (qualNameFromTypeName v)) substs ->
      applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs)
        `setUniqueness` u `addAliases` (<> als)
    | otherwise -> Scalar $ TypeVar als u v $ map substituteInTypeArg targs
  Scalar (Record ts) ->
    Scalar $ Record $ fmap (substituteTypes substs) ts
  Scalar (Arrow als v t1 t2) ->
    Scalar $ Arrow als v (substituteTypes substs t1) (substituteTypes substs t2)
  Scalar (Sum cs) ->
    Scalar $ Sum $ (fmap . fmap) (substituteTypes substs) cs
  where
    substituteInTypeArg (TypeArgDim d loc) =
      TypeArgDim (substituteInDim d) loc
    substituteInTypeArg (TypeArgType t loc) =
      TypeArgType (substituteTypes substs t) loc

    substituteInShape (ShapeDecl ds) =
      ShapeDecl $ map substituteInDim ds

    substituteInDim (NamedDim v)
      | Just (DimSub d) <- M.lookup (qualLeaf v) substs = d
    substituteInDim d = d

applyType ::
  Monoid als =>
  [TypeParam] ->
  TypeBase (DimDecl VName) als ->
  [StructTypeArg] ->
  TypeBase (DimDecl VName) als
applyType ps t args =
  substituteTypes substs t
  where
    substs = M.fromList $ zipWith mkSubst ps args
    -- We are assuming everything has already been type-checked for correctness.
    mkSubst (TypeParamDim pv _) (TypeArgDim (NamedDim v) _) =
      (pv, DimSub $ NamedDim v)
    mkSubst (TypeParamDim pv _) (TypeArgDim (ConstDim x) _) =
      (pv, DimSub $ ConstDim x)
    mkSubst (TypeParamDim pv _) (TypeArgDim AnyDim _) =
      (pv, DimSub AnyDim)
    mkSubst (TypeParamType l pv _) (TypeArgType at _) =
      (pv, TypeSub $ TypeAbbr l [] at)
    mkSubst p a =
      error $ "applyType mkSubst: cannot substitute " ++ pretty a ++ " for " ++ pretty p

-- | A type substituion may be a substitution or a yet-unknown
-- substitution (but which is certainly an overloaded primitive
-- type!).  The latter is used to remove aliases from types that are
-- yet-unknown but that we know cannot carry aliases (see issue #682).
data Subst t = Subst t | PrimSubst | SizeSubst (DimDecl VName)
  deriving (Show)

instance Functor Subst where
  fmap f (Subst t) = Subst $ f t
  fmap _ PrimSubst = PrimSubst
  fmap _ (SizeSubst v) = SizeSubst v

-- | Class of types which allow for substitution of types with no
-- annotations for type variable names.
class Substitutable a where
  applySubst :: (VName -> Maybe (Subst StructType)) -> a -> a

instance Substitutable (TypeBase (DimDecl VName) ()) where
  applySubst = substTypesAny

instance Substitutable (TypeBase (DimDecl VName) Aliasing) where
  applySubst = substTypesAny . (fmap (fmap fromStruct) .)

instance Substitutable (DimDecl VName) where
  applySubst f (NamedDim (QualName _ v))
    | Just (SizeSubst d) <- f v = d
  applySubst _ d = d

instance Substitutable d => Substitutable (ShapeDecl d) where
  applySubst f = fmap $ applySubst f

instance Substitutable Pattern where
  applySubst f = runIdentity . astMap mapper
    where
      mapper =
        ASTMapper
          { mapOnExp = return,
            mapOnName = return,
            mapOnQualName = return,
            mapOnStructType = return . applySubst f,
            mapOnPatternType = return . applySubst f
          }

-- | Perform substitutions, from type names to types, on a type. Works
-- regardless of what shape and uniqueness information is attached to the type.
substTypesAny ::
  Monoid as =>
  (VName -> Maybe (Subst (TypeBase (DimDecl VName) as))) ->
  TypeBase (DimDecl VName) as ->
  TypeBase (DimDecl VName) as
substTypesAny lookupSubst ot = case ot of
  Array als u et shape ->
    arrayOf
      (substTypesAny lookupSubst' (Scalar et))
      (applySubst lookupSubst' shape)
      u
      `setAliases` als
  Scalar (Prim t) -> Scalar $ Prim t
  -- We only substitute for a type variable with no arguments, since
  -- type parameters cannot have higher kind.
  Scalar (TypeVar als u v targs) ->
    case lookupSubst $ qualLeaf (qualNameFromTypeName v) of
      Just (Subst t) -> substTypesAny lookupSubst $ t `setUniqueness` u `addAliases` (<> als)
      Just PrimSubst -> Scalar $ TypeVar mempty u v $ map subsTypeArg targs
      _ -> Scalar $ TypeVar als u v $ map subsTypeArg targs
  Scalar (Record ts) -> Scalar $ Record $ fmap (substTypesAny lookupSubst) ts
  Scalar (Arrow als v t1 t2) ->
    Scalar $ Arrow als v (substTypesAny lookupSubst t1) (substTypesAny lookupSubst t2)
  Scalar (Sum ts) ->
    Scalar $ Sum $ fmap (fmap $ substTypesAny lookupSubst) ts
  where
    subsTypeArg (TypeArgType t loc) =
      TypeArgType (substTypesAny lookupSubst' t) loc
    subsTypeArg (TypeArgDim v loc) =
      TypeArgDim (applySubst lookupSubst' v) loc

    lookupSubst' = fmap (fmap $ second (const ())) . lookupSubst