{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | This module provides various simple ways to query and manipulate
-- fundamental Futhark terms, such as types and values.  The intent is to
-- keep "Futhark.Language.Syntax" simple, and put whatever embellishments
-- we need here.
module Language.Futhark.Prop
  ( -- * Various
    Intrinsic (..),
    intrinsics,
    maxIntrinsicTag,
    namesToPrimTypes,
    qualName,
    qualify,
    typeName,
    valueType,
    primValueType,
    leadingOperator,
    progImports,
    decImports,
    progModuleTypes,
    identifierReference,
    prettyStacktrace,

    -- * Queries on expressions
    typeOf,
    valBindTypeScheme,
    funType,

    -- * Queries on patterns and params
    patternIdents,
    patternNames,
    patternMap,
    patternType,
    patternStructType,
    patternParam,
    patternOrderZero,
    patternDimNames,

    -- * Queries on types
    uniqueness,
    unique,
    aliases,
    diet,
    arrayRank,
    arrayShape,
    nestedDims,
    orderZero,
    unfoldFunType,
    foldFunType,
    typeVars,
    typeDimNames,
    primByteSize,

    -- * Operations on types
    rank,
    peelArray,
    stripArray,
    arrayOf,
    toStructural,
    toStruct,
    fromStruct,
    setAliases,
    addAliases,
    setUniqueness,
    noSizes,
    anySizes,
    traverseDims,
    DimPos (..),
    mustBeExplicit,
    mustBeExplicitInType,
    tupleRecord,
    isTupleRecord,
    areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
    sortConstrs,
    isTypeParam,
    isSizeParam,
    combineTypeShapes,
    matchDims,
    unscopeType,
    onRecordField,
    -- | Values of these types are produces by the parser.  They use
    -- unadorned names and have no type information, apart from that
    -- which is syntactically required.
    NoInfo (..),
    UncheckedType,
    UncheckedTypeExp,
    UncheckedIdent,
    UncheckedTypeDecl,
    UncheckedDimIndex,
    UncheckedExp,
    UncheckedModExp,
    UncheckedSigExp,
    UncheckedTypeParam,
    UncheckedPattern,
    UncheckedValBind,
    UncheckedDec,
    UncheckedSpec,
    UncheckedProg,
    UncheckedCase,
  )
where

import Control.Monad.State
import Data.Bifoldable
import Data.Bifunctor
import Data.Bitraversable (bitraverse)
import Data.Char
import Data.Foldable
import Data.List (genericLength, isPrefixOf, sortOn)
import qualified Data.Map.Strict as M
import Data.Maybe
import Data.Ord
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Futhark.IR.Primitive as Primitive
import Futhark.Util (maxinum, nubOrd)
import Futhark.Util.Pretty
import Language.Futhark.Syntax

-- | Return the dimensionality of a type.  For non-arrays, this is
-- zero.  For a one-dimensional array it is one, for a two-dimensional
-- it is two, and so forth.
arrayRank :: TypeBase dim as -> Int
arrayRank = shapeRank . arrayShape

-- | Return the shape of a type - for non-arrays, this is 'mempty'.
arrayShape :: TypeBase dim as -> ShapeDecl dim
arrayShape (Array _ _ _ ds) = ds
arrayShape _ = mempty

-- | Return any shape declarations in the type, with duplicates
-- removed.
nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims t =
  case t of
    Array _ _ a ds ->
      nubOrd $ nestedDims (Scalar a) <> shapeDims ds
    Scalar (Record fs) ->
      nubOrd $ foldMap nestedDims fs
    Scalar Prim {} ->
      mempty
    Scalar (Sum cs) ->
      nubOrd $ foldMap (foldMap nestedDims) cs
    Scalar (Arrow _ v t1 t2) ->
      filter (notV v) $ nestedDims t1 <> nestedDims t2
    Scalar (TypeVar _ _ _ targs) ->
      concatMap typeArgDims targs
  where
    typeArgDims (TypeArgDim d _) = [d]
    typeArgDims (TypeArgType at _) = nestedDims at

    notV Unnamed = const True
    notV (Named v) = (/= NamedDim (qualName v))

-- | Change the shape of a type to be just the rank.
noSizes :: TypeBase (DimDecl vn) as -> TypeBase () as
noSizes = first $ const ()

-- | Change all size annotations to be 'AnyDim'.
anySizes :: TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
anySizes = first $ const AnyDim

-- | Where does this dimension occur?
data DimPos
  = -- | Immediately in the argument to 'traverseDims'.
    PosImmediate
  | -- | In a function parameter type.
    PosParam
  | -- | In a function return type.
    PosReturn
  deriving (Eq, Ord, Show)

-- | Perform a traversal (possibly including replacement) on sizes
-- that are parameters in a function type, but also including the type
-- immediately passed to the function.  Also passes along a set of the
-- parameter names inside the type that have come in scope at the
-- occurrence of the dimension.
traverseDims ::
  forall f fdim tdim als.
  Applicative f =>
  (S.Set VName -> DimPos -> fdim -> f tdim) ->
  TypeBase fdim als ->
  f (TypeBase tdim als)
traverseDims f = go mempty PosImmediate
  where
    go ::
      forall als'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim als' ->
      f (TypeBase tdim als')
    go bound b t@Array {} =
      bitraverse (f bound b) pure t
    go bound b (Scalar (Record fields)) =
      Scalar . Record <$> traverse (go bound b) fields
    go bound b (Scalar (TypeVar as u tn targs)) =
      Scalar <$> (TypeVar as u tn <$> traverse (onTypeArg bound b) targs)
    go bound b (Scalar (Sum cs)) =
      Scalar . Sum <$> traverse (traverse (go bound b)) cs
    go _ _ (Scalar (Prim t)) =
      pure $ Scalar $ Prim t
    go bound _ (Scalar (Arrow als p t1 t2)) =
      Scalar <$> (Arrow als p <$> go bound' PosParam t1 <*> go bound' PosReturn t2)
      where
        bound' = case p of
          Named p' -> S.insert p' bound
          Unnamed -> bound

    onTypeArg bound b (TypeArgDim d loc) =
      TypeArgDim <$> f bound b d <*> pure loc
    onTypeArg bound b (TypeArgType t loc) =
      TypeArgType <$> go bound b t <*> pure loc

mustBeExplicitAux :: StructType -> M.Map VName Bool
mustBeExplicitAux t =
  execState (traverseDims onDim t) mempty
  where
    onDim bound _ (NamedDim d)
      | qualLeaf d `S.member` bound =
        modify $ \s -> M.insertWith (&&) (qualLeaf d) False s
    onDim _ PosImmediate (NamedDim d) =
      modify $ \s -> M.insertWith (&&) (qualLeaf d) False s
    onDim _ _ (NamedDim d) =
      modify $ M.insertWith (&&) (qualLeaf d) True
    onDim _ _ _ =
      return ()

-- | Figure out which of the sizes in a parameter type must be passed
-- explicitly, because their first use is as something else than just
-- an array dimension.  'mustBeExplicit' is like this function, but
-- first decomposes into parameter types.
mustBeExplicitInType :: StructType -> S.Set VName
mustBeExplicitInType t =
  S.fromList $ M.keys $ M.filter id $ mustBeExplicitAux t

-- | Figure out which of the sizes in a binding type must be passed
-- explicitly, because their first use is as something else than just
-- an array dimension.
mustBeExplicit :: StructType -> S.Set VName
mustBeExplicit bind_t =
  let (ts, ret) = unfoldFunType bind_t
      alsoRet =
        M.unionWith (&&) $
          M.fromList $ zip (S.toList $ typeDimNames ret) $ repeat True
   in S.fromList $ M.keys $ M.filter id $ alsoRet $ foldl' onType mempty ts
  where
    onType uses t = uses <> mustBeExplicitAux t -- Left-biased union.

-- | Return the uniqueness of a type.
uniqueness :: TypeBase shape as -> Uniqueness
uniqueness (Array _ u _ _) = u
uniqueness (Scalar (TypeVar _ u _ _)) = u
uniqueness (Scalar (Sum ts)) = foldMap (foldMap uniqueness) $ M.elems ts
uniqueness (Scalar (Record fs)) = foldMap uniqueness $ M.elems fs
uniqueness _ = Nonunique

-- | @unique t@ is 'True' if the type of the argument is unique.
unique :: TypeBase shape as -> Bool
unique = (== Unique) . uniqueness

-- | Return the set of all variables mentioned in the aliasing of a
-- type.
aliases :: Monoid as => TypeBase shape as -> as
aliases = bifoldMap (const mempty) id

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ might consume its argument.
diet :: TypeBase shape as -> Diet
diet (Scalar (Record ets)) = RecordDiet $ fmap diet ets
diet (Scalar (Prim _)) = Observe
diet (Scalar (Arrow _ _ t1 t2)) = FuncDiet (diet t1) (diet t2)
diet (Array _ Unique _ _) = Consume
diet (Array _ Nonunique _ _) = Observe
diet (Scalar (TypeVar _ Unique _ _)) = Consume
diet (Scalar (TypeVar _ Nonunique _ _)) = Observe
diet (Scalar Sum {}) = Observe

-- | Convert any type to one that has rank information, no alias
-- information, and no embedded names.
toStructural ::
  TypeBase dim as ->
  TypeBase () ()
toStructural = flip setAliases () . first (const ())

-- | Remove aliasing information from a type.
toStruct ::
  TypeBase dim as ->
  TypeBase dim ()
toStruct t = t `setAliases` ()

-- | Replace no aliasing with an empty alias set.
fromStruct ::
  TypeBase dim as ->
  TypeBase dim Aliasing
fromStruct t = t `setAliases` S.empty

-- | @peelArray n t@ returns the type resulting from peeling the first
-- @n@ array dimensions from @t@.  Returns @Nothing@ if @t@ has less
-- than @n@ dimensions.
peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray n (Array als u t shape)
  | shapeRank shape == n =
    Just $ Scalar t `addAliases` const als
  | otherwise =
    Array als u t <$> stripDims n shape
peelArray _ _ = Nothing

-- | @arrayOf t s u@ constructs an array type.  The convenience
-- compared to using the 'Array' constructor directly is that @t@ can
-- itself be an array.  If @t@ is an @n@-dimensional array, and @s@ is
-- a list of length @n@, the resulting type is of an @n+m@ dimensions.
-- The uniqueness of the new array will be @u@, no matter the
-- uniqueness of @t@.
arrayOf ::
  Monoid as =>
  TypeBase dim as ->
  ShapeDecl dim ->
  Uniqueness ->
  TypeBase dim as
arrayOf t = arrayOfWithAliases (t `setUniqueness` Nonunique) mempty

arrayOfWithAliases ::
  Monoid as =>
  TypeBase dim as ->
  as ->
  ShapeDecl dim ->
  Uniqueness ->
  TypeBase dim as
arrayOfWithAliases (Array as1 _ et shape1) as2 shape2 u =
  Array (as1 <> as2) u et (shape2 <> shape1)
arrayOfWithAliases (Scalar t) as shape u =
  Array as u (second (const ()) t) shape

-- | @stripArray n t@ removes the @n@ outermost layers of the array.
-- Essentially, it is the type of indexing an array of type @t@ with
-- @n@ indexes.
stripArray :: Int -> TypeBase dim as -> TypeBase dim as
stripArray n (Array als u et shape)
  | Just shape' <- stripDims n shape =
    Array als u et shape'
  | otherwise =
    Scalar et `setUniqueness` u `setAliases` als
stripArray _ t = t

-- | Create a record type corresponding to a tuple with the given
-- element types.
tupleRecord :: [TypeBase dim as] -> TypeBase dim as
tupleRecord = Scalar . Record . M.fromList . zip tupleFieldNames

-- | Does this type corespond to a tuple?  If so, return the elements
-- of that tuple.
isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord (Scalar (Record fs)) = areTupleFields fs
isTupleRecord _ = Nothing

-- | Does this record map correspond to a tuple?
areTupleFields :: M.Map Name a -> Maybe [a]
areTupleFields fs =
  let fs' = sortFields fs
   in if and $ zipWith (==) (map fst fs') tupleFieldNames
        then Just $ map snd fs'
        else Nothing

-- | Construct a record map corresponding to a tuple.
tupleFields :: [a] -> M.Map Name a
tupleFields as = M.fromList $ zip tupleFieldNames as

-- | Increasing field names for a tuple (starts at 0).
tupleFieldNames :: [Name]
tupleFieldNames = map (nameFromString . show) [(0 :: Int) ..]

-- | Sort fields by their name; taking care to sort numeric fields by
-- their numeric value.  This ensures that tuples and tuple-like
-- records match.
sortFields :: M.Map Name a -> [(Name, a)]
sortFields l = map snd $ sortOn fst $ zip (map (fieldish . fst) l') l'
  where
    l' = M.toList l
    onDigit Nothing _ = Nothing
    onDigit (Just d) c
      | isDigit c = Just $ d * 10 + ord c - ord '0'
      | otherwise = Nothing
    fieldish s = maybe (Right s) Left $ T.foldl' onDigit (Just 0) $ nameToText s

-- | Sort the constructors of a sum type in some well-defined (but not
-- otherwise significant) manner.
sortConstrs :: M.Map Name a -> [(Name, a)]
sortConstrs cs = sortOn fst $ M.toList cs

-- | Is this a 'TypeParamType'?
isTypeParam :: TypeParamBase vn -> Bool
isTypeParam TypeParamType {} = True
isTypeParam TypeParamDim {} = False

-- | Is this a 'TypeParamDim'?
isSizeParam :: TypeParamBase vn -> Bool
isSizeParam = not . isTypeParam

-- | Combine the shape information of types as much as possible. The first
-- argument is the orignal type and the second is the type of the transformed
-- expression. This is necessary since the original type may contain additional
-- information (e.g., shape restrictions) from the user given annotation.
combineTypeShapes ::
  (Monoid as, ArrayDim dim) =>
  TypeBase dim as ->
  TypeBase dim as ->
  TypeBase dim as
combineTypeShapes (Scalar (Record ts1)) (Scalar (Record ts2))
  | M.keys ts1 == M.keys ts2 =
    Scalar $
      Record $
        M.map
          (uncurry combineTypeShapes)
          (M.intersectionWith (,) ts1 ts2)
combineTypeShapes (Scalar (Sum cs1)) (Scalar (Sum cs2))
  | M.keys cs1 == M.keys cs2 =
    Scalar $
      Sum $
        M.map
          (uncurry $ zipWith combineTypeShapes)
          (M.intersectionWith (,) cs1 cs2)
combineTypeShapes (Scalar (Arrow als1 p1 a1 b1)) (Scalar (Arrow als2 _p2 a2 b2)) =
  Scalar $ Arrow (als1 <> als2) p1 (combineTypeShapes a1 a2) (combineTypeShapes b1 b2)
combineTypeShapes (Scalar (TypeVar als1 u1 v targs1)) (Scalar (TypeVar als2 _ _ targs2)) =
  Scalar $ TypeVar (als1 <> als2) u1 v $ zipWith f targs1 targs2
  where
    f (TypeArgType t1 loc) (TypeArgType t2 _) =
      TypeArgType (combineTypeShapes t1 t2) loc
    f targ _ = targ
combineTypeShapes (Array als1 u1 et1 shape1) (Array als2 _u2 et2 _shape2) =
  arrayOfWithAliases
    ( combineTypeShapes (Scalar et1) (Scalar et2)
        `setAliases` mempty
    )
    (als1 <> als2)
    shape1
    u1
combineTypeShapes _ new_tp = new_tp

-- | Match the dimensions of otherwise assumed-equal types.
matchDims ::
  (Monoid as, Monad m) =>
  (d1 -> d2 -> m d1) ->
  TypeBase d1 as ->
  TypeBase d2 as ->
  m (TypeBase d1 as)
matchDims onDims t1 t2 =
  case (t1, t2) of
    (Array als1 u1 et1 shape1, Array als2 u2 et2 shape2) ->
      flip setAliases (als1 <> als2)
        <$> ( arrayOf
                <$> matchDims onDims (Scalar et1) (Scalar et2)
                <*> onShapes shape1 shape2
                <*> pure (min u1 u2)
            )
    (Scalar (Record f1), Scalar (Record f2)) ->
      Scalar . Record
        <$> traverse (uncurry (matchDims onDims)) (M.intersectionWith (,) f1 f2)
    (Scalar (Sum cs1), Scalar (Sum cs2)) ->
      Scalar . Sum
        <$> traverse
          (traverse (uncurry (matchDims onDims)))
          (M.intersectionWith zip cs1 cs2)
    (Scalar (Arrow als1 p1 a1 b1), Scalar (Arrow als2 _p2 a2 b2)) ->
      Scalar
        <$> (Arrow (als1 <> als2) p1 <$> matchDims onDims a1 a2 <*> matchDims onDims b1 b2)
    ( Scalar (TypeVar als1 u v targs1),
      Scalar (TypeVar als2 _ _ targs2)
      ) ->
        Scalar . TypeVar (als1 <> als2) u v <$> zipWithM matchTypeArg targs1 targs2
    _ -> return t1
  where
    matchTypeArg ta@TypeArgType {} _ = return ta
    matchTypeArg a _ = return a

    onShapes shape1 shape2 =
      ShapeDecl <$> zipWithM onDims (shapeDims shape1) (shapeDims shape2)

-- | Set the uniqueness attribute of a type.  If the type is a record
-- or sum type, the uniqueness of its components will be modified.
setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as
setUniqueness (Array als _ et shape) u =
  Array als u et shape
setUniqueness (Scalar (TypeVar als _ t targs)) u =
  Scalar $ TypeVar als u t targs
setUniqueness (Scalar (Record ets)) u =
  Scalar $ Record $ fmap (`setUniqueness` u) ets
setUniqueness (Scalar (Sum ets)) u =
  Scalar $ Sum $ fmap (map (`setUniqueness` u)) ets
setUniqueness t _ = t

-- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for
-- any already present aliasing.
setAliases :: TypeBase dim asf -> ast -> TypeBase dim ast
setAliases t = addAliases t . const

-- | @t \`addAliases\` f@ returns @t@, but with any already present
-- aliasing replaced by @f@ applied to that aliasing.
addAliases ::
  TypeBase dim asf ->
  (asf -> ast) ->
  TypeBase dim ast
addAliases = flip second

intValueType :: IntValue -> IntType
intValueType Int8Value {} = Int8
intValueType Int16Value {} = Int16
intValueType Int32Value {} = Int32
intValueType Int64Value {} = Int64

floatValueType :: FloatValue -> FloatType
floatValueType Float32Value {} = Float32
floatValueType Float64Value {} = Float64

-- | The type of a basic value.
primValueType :: PrimValue -> PrimType
primValueType (SignedValue v) = Signed $ intValueType v
primValueType (UnsignedValue v) = Unsigned $ intValueType v
primValueType (FloatValue v) = FloatType $ floatValueType v
primValueType BoolValue {} = Bool

-- | The type of the value.
valueType :: Value -> ValueType
valueType (PrimValue bv) = Scalar $ Prim $ primValueType bv
valueType (ArrayValue _ t) = t

-- | The size of values of this type, in bytes.
primByteSize :: Num a => PrimType -> a
primByteSize (Signed it) = Primitive.intByteSize it
primByteSize (Unsigned it) = Primitive.intByteSize it
primByteSize (FloatType ft) = Primitive.floatByteSize ft
primByteSize Bool = 1

-- | Construct a 'ShapeDecl' with the given number of 'AnyDim'
-- dimensions.
rank :: Int -> ShapeDecl (DimDecl VName)
rank n = ShapeDecl $ replicate n AnyDim

-- | The type is leaving a scope, so clean up any aliases that
-- reference the bound variables, and turn any dimensions that name
-- them into AnyDim instead.
unscopeType :: S.Set VName -> PatternType -> PatternType
unscopeType bound_here t = first onDim $ t `addAliases` S.map unbind
  where
    unbind (AliasBound v) | v `S.member` bound_here = AliasFree v
    unbind a = a
    onDim (NamedDim qn) | qualLeaf qn `S.member` bound_here = AnyDim
    onDim d = d

-- | Perform some operation on a given record field.  Returns
-- 'Nothing' if that field does not exist.
onRecordField ::
  (TypeBase dim als -> TypeBase dim als) ->
  [Name] ->
  TypeBase dim als ->
  Maybe (TypeBase dim als)
onRecordField f [] t = Just $ f t
onRecordField f (k : ks) (Scalar (Record m)) = do
  t <- onRecordField f ks =<< M.lookup k m
  Just $ Scalar $ Record $ M.insert k t m
onRecordField _ _ _ = Nothing

-- | The type of an Futhark term.  The aliasing will refer to itself, if
-- the term is a non-tuple-typed variable.
typeOf :: ExpBase Info VName -> PatternType
typeOf (Literal val _) = Scalar $ Prim $ primValueType val
typeOf (IntLit _ (Info t) _) = t
typeOf (FloatLit _ (Info t) _) = t
typeOf (Parens e _) = typeOf e
typeOf (QualParens _ e _) = typeOf e
typeOf (TupLit es _) = tupleRecord $ map typeOf es
typeOf (RecordLit fs _) =
  -- Reverse, because M.unions is biased to the left.
  Scalar $ Record $ M.unions $ reverse $ map record fs
  where
    record (RecordFieldExplicit name e _) = M.singleton name $ typeOf e
    record (RecordFieldImplicit name (Info t) _) =
      M.singleton (baseName name) $
        t
          `addAliases` S.insert (AliasBound name)
typeOf (ArrayLit _ (Info t) _) = t
typeOf (StringLit vs _) =
  Array
    mempty
    Unique
    (Prim (Unsigned Int8))
    (ShapeDecl [ConstDim $ genericLength vs])
typeOf (Range _ _ _ (Info t, _) _) = t
typeOf (BinOp _ _ _ _ (Info t) _ _) = t
typeOf (Project _ _ (Info t) _) = t
typeOf (If _ _ _ (Info t, _) _) = t
typeOf (Var _ (Info t) _) = t
typeOf (Ascript e _ _) = typeOf e
typeOf (Coerce _ _ (Info t, _) _) = t
typeOf (Apply _ _ _ (Info t, _) _) = t
typeOf (Negate e _) = typeOf e
typeOf (LetPat _ _ _ (Info t, _) _) = t
typeOf (LetFun _ _ _ (Info t) _) = t
typeOf (LetWith _ _ _ _ _ (Info t) _) = t
typeOf (Index _ _ (Info t, _) _) = t
typeOf (Update e _ _ _) = typeOf e `setAliases` mempty
typeOf (RecordUpdate _ _ _ (Info t) _) = t
typeOf (Assert _ e _ _) = typeOf e
typeOf (DoLoop _ _ _ _ _ (Info (t, _)) _) = t
typeOf (Lambda params _ _ (Info (als, t)) _) =
  unscopeType bound_here $ foldr (arrow . patternParam) t params `setAliases` als
  where
    bound_here =
      S.map identName (mconcat $ map patternIdents params)
        `S.difference` S.fromList (mapMaybe (named . patternParam) params)
    arrow (px, tx) y = Scalar $ Arrow () px tx y
    named (Named x, _) = Just x
    named (Unnamed, _) = Nothing
typeOf (OpSection _ (Info t) _) =
  t
typeOf (OpSectionLeft _ _ _ (_, Info (pn, pt2)) (Info ret, _) _) =
  Scalar $ Arrow mempty pn (fromStruct pt2) ret
typeOf (OpSectionRight _ _ _ (Info (pn, pt1), _) (Info ret) _) =
  Scalar $ Arrow mempty pn (fromStruct pt1) ret
typeOf (ProjectSection _ (Info t) _) = t
typeOf (IndexSection _ (Info t) _) = t
typeOf (Constr _ _ (Info t) _) = t
typeOf (Match _ cs (Info t, _) _) =
  unscopeType (foldMap unscopeSet cs) t
  where
    unscopeSet (CasePat p _ _) = S.map identName $ patternIdents p
typeOf (Attr _ e _) = typeOf e

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ps ret = foldr arrow ret ps
  where
    arrow t1 t2 = Scalar $ Arrow mempty Unnamed t1 t2

-- | Extract the parameter types and return type from a type.
-- If the type is not an arrow type, the list of parameter types is empty.
unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType (Scalar (Arrow _ _ t1 t2)) =
  let (ps, r) = unfoldFunType t2
   in (t1 : ps, r)
unfoldFunType t = ([], t)

-- | The type scheme of a value binding, comprising the type
-- parameters and the actual type.
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme vb =
  ( valBindTypeParams vb,
    funType (valBindParams vb) (fst (unInfo (valBindRetType vb)))
  )

-- | The type of a function with the given parameters and return type.
funType :: [PatternBase Info VName] -> StructType -> StructType
funType params ret = foldr (arrow . patternParam) ret params
  where
    arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt

-- | The type names mentioned in a type.
typeVars :: Monoid as => TypeBase dim as -> S.Set VName
typeVars t =
  case t of
    Scalar Prim {} -> mempty
    Scalar (TypeVar _ _ tn targs) ->
      mconcat $ typeVarFree tn : map typeArgFree targs
    Scalar (Arrow _ _ t1 t2) -> typeVars t1 <> typeVars t2
    Scalar (Record fields) -> foldMap typeVars fields
    Scalar (Sum cs) -> mconcat $ (foldMap . fmap) typeVars cs
    Array _ _ rt _ -> typeVars $ Scalar rt
  where
    typeVarFree = S.singleton . typeLeaf
    typeArgFree (TypeArgType ta _) = typeVars ta
    typeArgFree TypeArgDim {} = mempty

-- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not
-- a function type, does not contain a function type as a subcomponent, and may
-- not be instantiated with a function type.
orderZero :: TypeBase dim as -> Bool
orderZero Array {} = True
orderZero (Scalar (Prim _)) = True
orderZero (Scalar (Record fs)) = all orderZero $ M.elems fs
orderZero (Scalar TypeVar {}) = True
orderZero (Scalar Arrow {}) = False
orderZero (Scalar (Sum cs)) = all (all orderZero) cs

-- | Extract all the shape names that occur in a given pattern.
patternDimNames :: PatternBase Info VName -> S.Set VName
patternDimNames (TuplePattern ps _) = foldMap patternDimNames ps
patternDimNames (RecordPattern fs _) = foldMap (patternDimNames . snd) fs
patternDimNames (PatternParens p _) = patternDimNames p
patternDimNames (Id _ (Info tp) _) = typeDimNames tp
patternDimNames (Wildcard (Info tp) _) = typeDimNames tp
patternDimNames (PatternAscription p (TypeDecl _ (Info t)) _) =
  patternDimNames p <> typeDimNames t
patternDimNames (PatternLit _ (Info tp) _) = typeDimNames tp
patternDimNames (PatternConstr _ _ ps _) = foldMap patternDimNames ps

-- | Extract all the shape names that occur in a given type.
typeDimNames :: TypeBase (DimDecl VName) als -> S.Set VName
typeDimNames = foldMap dimName . nestedDims
  where
    dimName :: DimDecl VName -> S.Set VName
    dimName (NamedDim qn) = S.singleton $ qualLeaf qn
    dimName _ = mempty

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: PatternBase Info vn -> Bool
patternOrderZero pat = case pat of
  TuplePattern ps _ -> all patternOrderZero ps
  RecordPattern fs _ -> all (patternOrderZero . snd) fs
  PatternParens p _ -> patternOrderZero p
  Id _ (Info t) _ -> orderZero t
  Wildcard (Info t) _ -> orderZero t
  PatternAscription p _ _ -> patternOrderZero p
  PatternLit _ (Info t) _ -> orderZero t
  PatternConstr _ _ ps _ -> all patternOrderZero ps

-- | The set of identifiers bound in a pattern.
patternIdents :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn)
patternIdents (Id v t loc) = S.singleton $ Ident v t loc
patternIdents (PatternParens p _) = patternIdents p
patternIdents (TuplePattern pats _) = mconcat $ map patternIdents pats
patternIdents (RecordPattern fs _) = mconcat $ map (patternIdents . snd) fs
patternIdents Wildcard {} = mempty
patternIdents (PatternAscription p _ _) = patternIdents p
patternIdents PatternLit {} = mempty
patternIdents (PatternConstr _ _ ps _) = mconcat $ map patternIdents ps

-- | The set of names bound in a pattern.
patternNames :: (Functor f, Ord vn) => PatternBase f vn -> S.Set vn
patternNames (Id v _ _) = S.singleton v
patternNames (PatternParens p _) = patternNames p
patternNames (TuplePattern pats _) = mconcat $ map patternNames pats
patternNames (RecordPattern fs _) = mconcat $ map (patternNames . snd) fs
patternNames Wildcard {} = mempty
patternNames (PatternAscription p _ _) = patternNames p
patternNames PatternLit {} = mempty
patternNames (PatternConstr _ _ ps _) = mconcat $ map patternNames ps

-- | A mapping from names bound in a map to their identifier.
patternMap :: (Functor f) => PatternBase f VName -> M.Map VName (IdentBase f VName)
patternMap pat =
  M.fromList $ zip (map identName idents) idents
  where
    idents = S.toList $ patternIdents pat

-- | The type of values bound by the pattern.
patternType :: PatternBase Info VName -> PatternType
patternType (Wildcard (Info t) _) = t
patternType (PatternParens p _) = patternType p
patternType (Id _ (Info t) _) = t
patternType (TuplePattern pats _) = tupleRecord $ map patternType pats
patternType (RecordPattern fs _) = Scalar $ Record $ patternType <$> M.fromList fs
patternType (PatternAscription p _ _) = patternType p
patternType (PatternLit _ (Info t) _) = t
patternType (PatternConstr _ (Info t) _ _) = t

-- | The type matched by the pattern, including shape declarations if present.
patternStructType :: PatternBase Info VName -> StructType
patternStructType = toStruct . patternType

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: PatternBase Info VName -> (PName, StructType)
patternParam (PatternParens p _) =
  patternParam p
patternParam (PatternAscription (Id v _ _) td _) =
  (Named v, unInfo $ expandedType td)
patternParam (Id v (Info t) _) =
  (Named v, toStruct t)
patternParam p =
  (Unnamed, patternStructType p)

-- | Names of primitive types to types.  This is only valid if no
-- shadowing is going on, but useful for tools.
namesToPrimTypes :: M.Map Name PrimType
namesToPrimTypes =
  M.fromList
    [ (nameFromString $ pretty t, t)
      | t <-
          Bool :
          map Signed [minBound .. maxBound]
            ++ map Unsigned [minBound .. maxBound]
            ++ map FloatType [minBound .. maxBound]
    ]

-- | The nature of something predefined.  These can either be
-- monomorphic or overloaded.  An overloaded builtin is a list valid
-- types it can be instantiated with, to the parameter and result
-- type, with 'Nothing' representing the overloaded parameter type.
data Intrinsic
  = IntrinsicMonoFun [PrimType] PrimType
  | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType)
  | IntrinsicPolyFun [TypeParamBase VName] [StructType] StructType
  | IntrinsicType PrimType
  | IntrinsicEquality -- Special cased.

-- | A map of all built-ins.
intrinsics :: M.Map VName Intrinsic
intrinsics =
  M.fromList $
    zipWith namify [10 ..] $
      map primFun (M.toList Primitive.primFuns)
        ++ [("opaque", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)]
        ++ map unOpFun Primitive.allUnOps
        ++ map binOpFun Primitive.allBinOps
        ++ map cmpOpFun Primitive.allCmpOps
        ++ map convOpFun Primitive.allConvOps
        ++ map signFun Primitive.allIntTypes
        ++ map unsignFun Primitive.allIntTypes
        ++ map
          intrinsicType
          ( map Signed [minBound .. maxBound]
              ++ map Unsigned [minBound .. maxBound]
              ++ map FloatType [minBound .. maxBound]
              ++ [Bool]
          )
        ++
        -- This overrides the ! from Primitive.
        [ ( "!",
            IntrinsicOverloadedFun
              ( map Signed [minBound .. maxBound]
                  ++ map Unsigned [minBound .. maxBound]
                  ++ [Bool]
              )
              [Nothing]
              Nothing
          )
        ]
        ++
        -- The reason for the loop formulation is to ensure that we
        -- get a missing case warning if we forget a case.
        mapMaybe mkIntrinsicBinOp [minBound .. maxBound]
        ++ [ ( "flatten",
               IntrinsicPolyFun
                 [tp_a]
                 [Array () Nonunique t_a (rank 2)]
                 $ Array () Nonunique t_a (rank 1)
             ),
             ( "unflatten",
               IntrinsicPolyFun
                 [tp_a]
                 [ Scalar $ Prim $ Signed Int64,
                   Scalar $ Prim $ Signed Int64,
                   Array () Nonunique t_a (rank 1)
                 ]
                 $ Array () Nonunique t_a (rank 2)
             ),
             ( "concat",
               IntrinsicPolyFun
                 [tp_a]
                 [arr_a, arr_a]
                 uarr_a
             ),
             ( "rotate",
               IntrinsicPolyFun
                 [tp_a]
                 [Scalar $ Prim $ Signed Int64, arr_a]
                 arr_a
             ),
             ("transpose", IntrinsicPolyFun [tp_a] [arr_2d_a] arr_2d_a),
             ( "scatter",
               IntrinsicPolyFun
                 [tp_a]
                 [ Array () Unique t_a (rank 1),
                   Array () Nonunique (Prim $ Signed Int64) (rank 1),
                   Array () Nonunique t_a (rank 1)
                 ]
                 $ Array () Unique t_a (rank 1)
             ),
             ( "scatter_2d",
               IntrinsicPolyFun
                 [tp_a]
                 [ uarr_2d_a,
                   Array () Nonunique (tupInt64 2) (rank 1),
                   Array () Nonunique t_a (rank 1)
                 ]
                 uarr_2d_a
             ),
             ( "scatter_3d",
               IntrinsicPolyFun
                 [tp_a]
                 [ uarr_3d_a,
                   Array () Nonunique (tupInt64 3) (rank 1),
                   Array () Nonunique t_a (rank 1)
                 ]
                 uarr_3d_a
             ),
             ("zip", IntrinsicPolyFun [tp_a, tp_b] [arr_a, arr_b] uarr_a_b),
             ("unzip", IntrinsicPolyFun [tp_a, tp_b] [arr_a_b] t_arr_a_arr_b),
             ( "hist",
               IntrinsicPolyFun
                 [tp_a]
                 [ Scalar $ Prim $ Signed Int64,
                   uarr_a,
                   Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),
                   Scalar t_a,
                   Array () Nonunique (Prim $ Signed Int64) (rank 1),
                   arr_a
                 ]
                 uarr_a
             ),
             ("map", IntrinsicPolyFun [tp_a, tp_b] [Scalar t_a `arr` Scalar t_b, arr_a] uarr_b),
             ( "reduce",
               IntrinsicPolyFun
                 [tp_a]
                 [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a]
                 $ Scalar t_a
             ),
             ( "reduce_comm",
               IntrinsicPolyFun
                 [tp_a]
                 [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a]
                 $ Scalar t_a
             ),
             ( "scan",
               IntrinsicPolyFun
                 [tp_a]
                 [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a]
                 uarr_a
             ),
             ( "partition",
               IntrinsicPolyFun
                 [tp_a]
                 [ Scalar (Prim $ Signed Int32),
                   Scalar t_a `arr` Scalar (Prim $ Signed Int64),
                   arr_a
                 ]
                 $ tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int64) (rank 1)]
             ),
             ( "map_stream",
               IntrinsicPolyFun
                 [tp_a, tp_b]
                 [Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), arr_a]
                 uarr_b
             ),
             ( "map_stream_per",
               IntrinsicPolyFun
                 [tp_a, tp_b]
                 [Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), arr_a]
                 uarr_b
             ),
             ( "reduce_stream",
               IntrinsicPolyFun
                 [tp_a, tp_b]
                 [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b),
                   Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b),
                   arr_a
                 ]
                 $ Scalar t_b
             ),
             ( "reduce_stream_per",
               IntrinsicPolyFun
                 [tp_a, tp_b]
                 [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b),
                   Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b),
                   arr_a
                 ]
                 $ Scalar t_b
             ),
             ("trace", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a),
             ("break", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)
           ]
  where
    tv_a = VName (nameFromString "a") 0
    t_a = TypeVar () Nonunique (typeName tv_a) []
    arr_a = Array () Nonunique t_a (rank 1)
    arr_2d_a = Array () Nonunique t_a (rank 2)
    uarr_2d_a = Array () Unique t_a (rank 2)
    uarr_3d_a = Array () Unique t_a (rank 3)
    uarr_a = Array () Unique t_a (rank 1)
    tp_a = TypeParamType Unlifted tv_a mempty

    tv_b = VName (nameFromString "b") 1
    t_b = TypeVar () Nonunique (typeName tv_b) []
    arr_b = Array () Nonunique t_b (rank 1)
    uarr_b = Array () Unique t_b (rank 1)
    tp_b = TypeParamType Unlifted tv_b mempty

    arr_a_b =
      Array
        ()
        Nonunique
        (Record (M.fromList $ zip tupleFieldNames [Scalar t_a, Scalar t_b]))
        (rank 1)
    uarr_a_b = arr_a_b `setUniqueness` Unique
    t_arr_a_arr_b = Scalar $ Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b]

    arr x y = Scalar $ Arrow mempty Unnamed x y

    kv = VName (nameFromString "k") 2
    arr_ka = Array () Nonunique t_a (ShapeDecl [NamedDim $ qualName kv])
    arr_kb = Array () Nonunique t_b (ShapeDecl [NamedDim $ qualName kv])
    karr x y = Scalar $ Arrow mempty (Named kv) x y

    namify i (k, v) = (VName (nameFromString k) i, v)

    primFun (name, (ts, t, _)) =
      (name, IntrinsicMonoFun (map unPrim ts) $ unPrim t)

    unOpFun bop = (pretty bop, IntrinsicMonoFun [t] t)
      where
        t = unPrim $ Primitive.unOpType bop

    binOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] t)
      where
        t = unPrim $ Primitive.binOpType bop

    cmpOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] Bool)
      where
        t = unPrim $ Primitive.cmpOpType bop

    convOpFun cop = (pretty cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt)
      where
        (ft, tt) = Primitive.convOpType cop

    signFun t = ("sign_" ++ pretty t, IntrinsicMonoFun [Unsigned t] $ Signed t)

    unsignFun t = ("unsign_" ++ pretty t, IntrinsicMonoFun [Signed t] $ Unsigned t)

    unPrim (Primitive.IntType t) = Signed t
    unPrim (Primitive.FloatType t) = FloatType t
    unPrim Primitive.Bool = Bool
    unPrim Primitive.Cert = Bool

    intrinsicType t = (pretty t, IntrinsicType t)

    anyIntType =
      map Signed [minBound .. maxBound]
        ++ map Unsigned [minBound .. maxBound]
    anyNumberType =
      anyIntType
        ++ map FloatType [minBound .. maxBound]
    anyPrimType = Bool : anyNumberType

    mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic)
    mkIntrinsicBinOp op = do
      op' <- intrinsicBinOp op
      return (pretty op, op')

    binOp ts = Just $ IntrinsicOverloadedFun ts [Nothing, Nothing] Nothing
    ordering = Just $ IntrinsicOverloadedFun anyPrimType [Nothing, Nothing] (Just Bool)

    intrinsicBinOp Plus = binOp anyNumberType
    intrinsicBinOp Minus = binOp anyNumberType
    intrinsicBinOp Pow = binOp anyNumberType
    intrinsicBinOp Times = binOp anyNumberType
    intrinsicBinOp Divide = binOp anyNumberType
    intrinsicBinOp Mod = binOp anyNumberType
    intrinsicBinOp Quot = binOp anyIntType
    intrinsicBinOp Rem = binOp anyIntType
    intrinsicBinOp ShiftR = binOp anyIntType
    intrinsicBinOp ShiftL = binOp anyIntType
    intrinsicBinOp Band = binOp anyIntType
    intrinsicBinOp Xor = binOp anyIntType
    intrinsicBinOp Bor = binOp anyIntType
    intrinsicBinOp LogAnd = binOp [Bool]
    intrinsicBinOp LogOr = binOp [Bool]
    intrinsicBinOp Equal = Just IntrinsicEquality
    intrinsicBinOp NotEqual = Just IntrinsicEquality
    intrinsicBinOp Less = ordering
    intrinsicBinOp Leq = ordering
    intrinsicBinOp Greater = ordering
    intrinsicBinOp Geq = ordering
    intrinsicBinOp _ = Nothing

    tupInt64 n =
      Record $
        M.fromList $
          zip tupleFieldNames $
            replicate n $ Scalar $ Prim $ Signed Int64

-- | The largest tag used by an intrinsic - this can be used to
-- determine whether a 'VName' refers to an intrinsic or a user-defined name.
maxIntrinsicTag :: Int
maxIntrinsicTag = maxinum $ map baseTag $ M.keys intrinsics

-- | Create a name with no qualifiers from a name.
qualName :: v -> QualName v
qualName = QualName []

-- | Add another qualifier (at the head) to a qualified name.
qualify :: v -> QualName v -> QualName v
qualify k (QualName ks v) = QualName (k : ks) v

-- | Create a type name name with no qualifiers from a 'VName'.
typeName :: VName -> TypeName
typeName = typeNameFromQualName . qualName

-- | The modules imported by a Futhark program.
progImports :: ProgBase f vn -> [(String, SrcLoc)]
progImports = concatMap decImports . progDecs

-- | The modules imported by a single declaration.
decImports :: DecBase f vn -> [(String, SrcLoc)]
decImports (OpenDec x _) = modExpImports x
decImports (ModDec md) = modExpImports $ modExp md
decImports SigDec {} = []
decImports TypeDec {} = []
decImports ValDec {} = []
decImports (LocalDec d _) = decImports d
decImports (ImportDec x _ loc) = [(x, loc)]

modExpImports :: ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModVar {} = []
modExpImports (ModParens p _) = modExpImports p
modExpImports (ModImport f _ loc) = [(f, loc)]
modExpImports (ModDecs ds _) = concatMap decImports ds
modExpImports (ModApply _ me _ _ _) = modExpImports me
modExpImports (ModAscript me _ _ _) = modExpImports me
modExpImports ModLambda {} = []

-- | The set of module types used in any exported (non-local)
-- declaration.
progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn
progModuleTypes = mconcat . map onDec . progDecs
  where
    onDec (OpenDec x _) = onModExp x
    onDec (ModDec md) =
      maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md)
    onDec SigDec {} = mempty
    onDec TypeDec {} = mempty
    onDec ValDec {} = mempty
    onDec LocalDec {} = mempty
    onDec ImportDec {} = mempty

    onModExp ModVar {} = mempty
    onModExp (ModParens p _) = onModExp p
    onModExp ModImport {} = mempty
    onModExp (ModDecs ds _) = mconcat $ map onDec ds
    onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2
    onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se
    onModExp (ModLambda p r me _) =
      onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me

    onModParam = onSigExp . modParamType

    onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v
    onSigExp (SigParens e _) = onSigExp e
    onSigExp SigSpecs {} = mempty
    onSigExp (SigWith e _ _) = onSigExp e
    onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2

-- | Extract a leading @((name, namespace, file), remainder)@ from a
-- documentation comment string.  These are formatted as
-- \`name\`\@namespace[\@file].  Let us hope that this pattern does not occur
-- anywhere else.
identifierReference :: String -> Maybe ((String, String, Maybe FilePath), String)
identifierReference ('`' : s)
  | (identifier, '`' : '@' : s') <- break (== '`') s,
    (namespace, s'') <- span isAlpha s',
    not $ null namespace =
    case s'' of
      '@' : '"' : s'''
        | (file, '"' : s'''') <- span (/= '"') s''' ->
          Just ((identifier, namespace, Just file), s'''')
      _ -> Just ((identifier, namespace, Nothing), s'')
identifierReference _ = Nothing

-- | Given an operator name, return the operator that determines its
-- syntactical properties.
leadingOperator :: Name -> BinOp
leadingOperator s =
  maybe Backtick snd $
    find ((`isPrefixOf` s') . fst) $
      sortOn (Down . length . fst) $
        zip (map pretty operators) operators
  where
    s' = nameToString s
    operators :: [BinOp]
    operators = [minBound .. maxBound :: BinOp]

-- | A type with no aliasing information but shape annotations.
type UncheckedType = TypeBase (ShapeDecl Name) ()

-- | An expression with no type annotations.
type UncheckedTypeExp = TypeExp Name

-- | A type declaration with no expanded type.
type UncheckedTypeDecl = TypeDeclBase NoInfo Name

-- | An identifier with no type annotations.
type UncheckedIdent = IdentBase NoInfo Name

-- | An index with no type annotations.
type UncheckedDimIndex = DimIndexBase NoInfo Name

-- | An expression with no type annotations.
type UncheckedExp = ExpBase NoInfo Name

-- | A module expression with no type annotations.
type UncheckedModExp = ModExpBase NoInfo Name

-- | A module type expression with no type annotations.
type UncheckedSigExp = SigExpBase NoInfo Name

-- | A type parameter with no type annotations.
type UncheckedTypeParam = TypeParamBase Name

-- | A pattern with no type annotations.
type UncheckedPattern = PatternBase NoInfo Name

-- | A function declaration with no type annotations.
type UncheckedValBind = ValBindBase NoInfo Name

-- | A declaration with no type annotations.
type UncheckedDec = DecBase NoInfo Name

-- | A spec with no type annotations.
type UncheckedSpec = SpecBase NoInfo Name

-- | A Futhark program with no type annotations.
type UncheckedProg = ProgBase NoInfo Name

-- | A case (of a match expression) with no type annotations.
type UncheckedCase = CaseBase NoInfo Name