-- | 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,
    intrinsicVar,
    isBuiltin,
    isBuiltinLoc,
    maxIntrinsicTag,
    namesToPrimTypes,
    qualName,
    qualify,
    primValueType,
    leadingOperator,
    progImports,
    decImports,
    progModuleTypes,
    identifierReference,
    prettyStacktrace,
    progHoles,
    defaultEntryPoint,
    paramName,
    anySize,

    -- * Queries on expressions
    typeOf,
    valBindTypeScheme,
    valBindBound,
    funType,
    stripExp,
    similarExps,

    -- * Queries on patterns and params
    patIdents,
    patNames,
    patternMap,
    patternType,
    patternStructType,
    patternParam,
    patternOrderZero,

    -- * Queries on types
    uniqueness,
    unique,
    diet,
    arrayRank,
    arrayShape,
    orderZero,
    unfoldFunType,
    foldFunType,
    typeVars,
    isAccType,

    -- * Operations on types
    peelArray,
    stripArray,
    arrayOf,
    arrayOfWithAliases,
    toStructural,
    toStruct,
    toRes,
    toParam,
    resToParam,
    paramToRes,
    toResRet,
    setUniqueness,
    noSizes,
    traverseDims,
    DimPos (..),
    tupleRecord,
    isTupleRecord,
    areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
    sortConstrs,
    isTypeParam,
    isSizeParam,
    matchDims,

    -- * Un-typechecked ASTs
    UncheckedType,
    UncheckedTypeExp,
    UncheckedIdent,
    UncheckedDimIndex,
    UncheckedSlice,
    UncheckedExp,
    UncheckedModExp,
    UncheckedModTypeExp,
    UncheckedTypeParam,
    UncheckedPat,
    UncheckedValBind,
    UncheckedTypeBind,
    UncheckedModTypeBind,
    UncheckedModBind,
    UncheckedDec,
    UncheckedSpec,
    UncheckedProg,
    UncheckedCase,

    -- * Type-checked ASTs
    Ident,
    DimIndex,
    Slice,
    AppExp,
    Exp,
    Pat,
    ModExp,
    ModParam,
    ModTypeExp,
    ModBind,
    ModTypeBind,
    ValBind,
    Dec,
    Spec,
    Prog,
    TypeBind,
    StructTypeArg,
    ScalarType,
    TypeParam,
    Case,
  )
where

import Control.Monad
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable (bitraverse)
import Data.Char
import Data.Foldable
import Data.List (genericLength, isPrefixOf, sortOn)
import Data.List.NonEmpty qualified as NE
import Data.Loc (Loc (..), posFile)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Data.Set qualified as S
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import Language.Futhark.Primitive qualified as Primitive
import Language.Futhark.Syntax
import Language.Futhark.Traversals
import Language.Futhark.Tuple
import System.FilePath (takeDirectory)

-- | The name of the default program entry point (@main@).
defaultEntryPoint :: Name
defaultEntryPoint = nameFromString "main"

-- | 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 -> Shape dim
arrayShape (Array _ ds _) = ds
arrayShape _ = mempty

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

-- | 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 tn targs)) =
      Scalar <$> (TypeVar as tn <$> traverse (onTypeArg tn 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 u t1 (RetType dims t2))) =
      Scalar <$> (Arrow als p u <$> go bound' PosParam t1 <*> (RetType dims <$> go bound' PosReturn t2))
      where
        bound' =
          S.fromList dims
            <> case p of
              Named p' -> S.insert p' bound
              Unnamed -> bound

    onTypeArg _ bound b (TypeArgDim d) =
      TypeArgDim <$> f bound b d
    onTypeArg tn bound b (TypeArgType t) =
      TypeArgType <$> go bound b' t
      where
        b' =
          if qualLeaf tn == fst intrinsicAcc
            then b
            else PosParam

-- | Return the uniqueness of a type.
uniqueness :: TypeBase shape Uniqueness -> Uniqueness
uniqueness (Array u _ _) = u
uniqueness (Scalar (TypeVar u _ _)) = u
uniqueness (Scalar (Sum ts))
  | any (any unique) ts = Unique
uniqueness (Scalar (Record fs))
  | any unique fs = Unique
uniqueness _ = Nonunique

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

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ consumes its argument.
diet :: TypeBase shape Diet -> Diet
diet (Scalar (Record ets)) = foldl max Observe $ fmap diet ets
diet (Scalar (Prim _)) = Observe
diet (Scalar (Arrow {})) = Observe
diet (Array d _ _) = d
diet (Scalar (TypeVar d _ _)) = d
diet (Scalar (Sum cs)) = foldl max Observe $ foldMap (map diet) cs

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

-- | Remove uniquenss information from a type.
toStruct :: TypeBase dim u -> TypeBase dim NoUniqueness
toStruct = second (const NoUniqueness)

-- | Uses 'Observe'.
toParam :: Diet -> TypeBase Size u -> ParamType
toParam d = fmap (const d)

-- | Convert to 'ResType'
toRes :: Uniqueness -> TypeBase Size u -> ResType
toRes u = fmap (const u)

-- | Convert to 'ResRetType'
toResRet :: Uniqueness -> RetTypeBase Size u -> ResRetType
toResRet u = second (const u)

-- | Preserves relation between 'Diet' and 'Uniqueness'.
resToParam :: ResType -> ParamType
resToParam = second f
  where
    f Unique = Consume
    f Nonunique = Observe

-- | Preserves relation between 'Diet' and 'Uniqueness'.
paramToRes :: ParamType -> ResType
paramToRes = second f
  where
    f Consume = Unique
    f Observe = Nonunique

-- | @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 u -> Maybe (TypeBase dim u)
peelArray n (Array u shape t)
  | shapeRank shape == n =
      Just $ second (const u) (Scalar t)
  | otherwise =
      Array u <$> stripDims n shape <*> pure t
peelArray _ _ = Nothing

-- | @arrayOf u s t@ 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.
arrayOf ::
  Shape dim ->
  TypeBase dim NoUniqueness ->
  TypeBase dim NoUniqueness
arrayOf = arrayOfWithAliases mempty

-- | Like 'arrayOf', but you can pass in uniqueness info of the
-- resulting array.
arrayOfWithAliases ::
  u ->
  Shape dim ->
  TypeBase dim u ->
  TypeBase dim u
arrayOfWithAliases u shape2 (Array _ shape1 et) =
  Array u (shape2 <> shape1) et
arrayOfWithAliases u shape (Scalar t) =
  Array u shape (second (const mempty) t)

-- | @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 u shape et)
  | Just shape' <- stripDims n shape =
      Array u shape' et
  | otherwise =
      second (const u) (Scalar et)
stripArray _ t = t

-- | Create a record type corresponding to a tuple with the given
-- element types.
tupleRecord :: [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord = 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

-- | 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

-- | The name, if any.
paramName :: PName -> Maybe VName
paramName (Named v) = Just v
paramName Unnamed = Nothing

-- | A special expression representing no known size.  When present in
-- a type, each instance represents a distinct size.  The type checker
-- should _never_ produce these - they are a (hopefully temporary)
-- thing introduced by defunctorisation and monomorphisation.  They
-- represent a flaw in our implementation.  When they occur in a
-- return type, they can be replaced with freshly created existential
-- sizes.  When they occur in parameter types, they can be replaced
-- with size parameters.
anySize :: Size
anySize =
  -- The definition here is weird to avoid seeing this as a free
  -- variable.
  StringLit [65, 78, 89] mempty

-- | Match the dimensions of otherwise assumed-equal types.  The
-- combining function is also passed the names bound within the type
-- (from named parameters or return types).
matchDims ::
  forall as m d1 d2.
  (Monoid as, Monad m) =>
  ([VName] -> d1 -> d2 -> m d1) ->
  TypeBase d1 as ->
  TypeBase d2 as ->
  m (TypeBase d1 as)
matchDims onDims = matchDims' mempty
  where
    matchDims' ::
      forall u'. (Monoid u') => [VName] -> TypeBase d1 u' -> TypeBase d2 u' -> m (TypeBase d1 u')
    matchDims' bound t1 t2 =
      case (t1, t2) of
        (Array u1 shape1 et1, Array u2 shape2 et2) ->
          arrayOfWithAliases u1
            <$> onShapes bound shape1 shape2
            <*> matchDims' bound (second (const u2) (Scalar et1)) (second (const u2) (Scalar et2))
        (Scalar (Record f1), Scalar (Record f2)) ->
          Scalar . Record
            <$> traverse (uncurry (matchDims' bound)) (M.intersectionWith (,) f1 f2)
        (Scalar (Sum cs1), Scalar (Sum cs2)) ->
          Scalar . Sum
            <$> traverse
              (traverse (uncurry (matchDims' bound)))
              (M.intersectionWith zip cs1 cs2)
        ( Scalar (Arrow als1 p1 d1 a1 (RetType dims1 b1)),
          Scalar (Arrow als2 p2 _d2 a2 (RetType dims2 b2))
          ) ->
            let bound' = mapMaybe paramName [p1, p2] <> dims1 <> dims2 <> bound
             in Scalar
                  <$> ( Arrow (als1 <> als2) p1 d1
                          <$> matchDims' bound' a1 a2
                          <*> (RetType dims1 <$> matchDims' bound' b1 b2)
                      )
        ( Scalar (TypeVar als1 v targs1),
          Scalar (TypeVar als2 _ targs2)
          ) ->
            Scalar . TypeVar (als1 <> als2) v
              <$> zipWithM (matchTypeArg bound) targs1 targs2
        _ -> pure t1

    matchTypeArg bound (TypeArgType t1) (TypeArgType t2) =
      TypeArgType <$> matchDims' bound t1 t2
    matchTypeArg bound (TypeArgDim x) (TypeArgDim y) =
      TypeArgDim <$> onDims bound x y
    matchTypeArg _ a _ = pure a

    onShapes bound shape1 shape2 =
      Shape <$> zipWithM (onDims bound) (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 u1 -> u2 -> TypeBase dim u2
setUniqueness t u = second (const u) t

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

floatValueType :: FloatValue -> FloatType
floatValueType Float16Value {} = Float16
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 an Futhark term.  The aliasing will refer to itself, if
-- the term is a non-tuple-typed variable.
typeOf :: ExpBase Info VName -> StructType
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 _) = Scalar $ tupleRecord $ map typeOf es
typeOf (RecordLit fs _) =
  Scalar $ Record $ M.fromList $ map record fs
  where
    record (RecordFieldExplicit name e _) = (name, typeOf e)
    record (RecordFieldImplicit name (Info t) _) = (baseName name, t)
typeOf (ArrayLit _ (Info t) _) = t
typeOf (StringLit vs loc) =
  Array
    mempty
    (Shape [sizeFromInteger (genericLength vs) loc])
    (Prim (Unsigned Int8))
typeOf (Project _ _ (Info t) _) = t
typeOf (Var _ (Info t) _) = t
typeOf (Hole (Info t) _) = t
typeOf (Ascript e _ _) = typeOf e
typeOf (Coerce _ _ (Info t) _) = t
typeOf (Negate e _) = typeOf e
typeOf (Not e _) = typeOf e
typeOf (Update e _ _ _) = typeOf e
typeOf (RecordUpdate _ _ _ (Info t) _) = t
typeOf (Assert _ e _ _) = typeOf e
typeOf (Lambda params _ _ (Info t) _) = funType params t
typeOf (OpSection _ (Info t) _) = t
typeOf (OpSectionLeft _ _ _ (_, Info (pn, pt2)) (Info ret, _) _) =
  Scalar $ Arrow mempty pn (diet pt2) (toStruct pt2) ret
typeOf (OpSectionRight _ _ _ (Info (pn, pt1), _) (Info ret) _) =
  Scalar $ Arrow mempty pn (diet pt1) (toStruct pt1) ret
typeOf (ProjectSection _ (Info t) _) = t
typeOf (IndexSection _ (Info t) _) = t
typeOf (Constr _ _ (Info t) _) = t
typeOf (Attr _ e _) = typeOf e
typeOf (AppExp _ (Info res)) = appResType res

-- | The type of a function with the given parameters and return type.
funType :: [Pat ParamType] -> ResRetType -> StructType
funType params ret =
  let RetType _ t = foldr (arrow . patternParam) ret params
   in toStruct t
  where
    arrow (xp, d, xt) yt =
      RetType [] $ Scalar $ Arrow Nonunique xp d xt yt

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType :: [ParamType] -> ResRetType -> StructType
foldFunType ps ret =
  let RetType _ t = foldr arrow ret ps
   in toStruct t
  where
    arrow t1 t2 =
      RetType [] $ Scalar $ Arrow Nonunique Unnamed (diet t1) (toStruct 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 Diet], TypeBase dim NoUniqueness)
unfoldFunType (Scalar (Arrow _ _ d t1 (RetType _ t2))) =
  let (ps, r) = unfoldFunType t2
   in (second (const d) t1 : ps, r)
unfoldFunType t = ([], toStruct 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) (unInfo (valBindRetType vb))
  )

-- | The names that are brought into scope by this value binding (not
-- including its own parameter names, but including any existential
-- sizes).
valBindBound :: ValBindBase Info VName -> [VName]
valBindBound vb =
  valBindName vb
    : case valBindParams vb of
      [] -> retDims (unInfo (valBindRetType vb))
      _ -> []

-- | The type names mentioned in a type.
typeVars :: TypeBase dim as -> S.Set VName
typeVars t =
  case t of
    Scalar Prim {} -> mempty
    Scalar (TypeVar _ tn targs) ->
      mconcat $ S.singleton (qualLeaf tn) : map typeArgFree targs
    Scalar (Arrow _ _ _ t1 (RetType _ 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
    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

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: Pat (TypeBase d u) -> Bool
patternOrderZero = orderZero . patternType

-- | The set of identifiers bound in a pattern.
patIdents :: Pat t -> [Ident t]
patIdents (Id v t loc) = [Ident v t loc]
patIdents (PatParens p _) = patIdents p
patIdents (TuplePat pats _) = foldMap patIdents pats
patIdents (RecordPat fs _) = foldMap (patIdents . snd) fs
patIdents Wildcard {} = mempty
patIdents (PatAscription p _ _) = patIdents p
patIdents PatLit {} = mempty
patIdents (PatConstr _ _ ps _) = foldMap patIdents ps
patIdents (PatAttr _ p _) = patIdents p

-- | The set of names bound in a pattern.
patNames :: Pat t -> [VName]
patNames = map fst . patternMap

-- | Each name bound in a pattern alongside its type.
patternMap :: Pat t -> [(VName, t)]
patternMap = map f . patIdents
  where
    f (Ident v (Info t) _) = (v, t)

-- | The type of values bound by the pattern.
patternType :: Pat (TypeBase d u) -> TypeBase d u
patternType (Wildcard (Info t) _) = t
patternType (PatParens p _) = patternType p
patternType (Id _ (Info t) _) = t
patternType (TuplePat pats _) = Scalar $ tupleRecord $ map patternType pats
patternType (RecordPat fs _) = Scalar $ Record $ patternType <$> M.fromList fs
patternType (PatAscription p _ _) = patternType p
patternType (PatLit _ (Info t) _) = t
patternType (PatConstr _ (Info t) _ _) = t
patternType (PatAttr _ p _) = patternType p

-- | The type matched by the pattern, including shape declarations if present.
patternStructType :: Pat (TypeBase Size u) -> StructType
patternStructType = toStruct . patternType

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: Pat ParamType -> (PName, Diet, StructType)
patternParam (PatParens p _) =
  patternParam p
patternParam (PatAttr _ p _) =
  patternParam p
patternParam (PatAscription (Id v (Info t) _) _ _) =
  (Named v, diet t, toStruct t)
patternParam (Id v (Info t) _) =
  (Named v, diet t, toStruct t)
patternParam p =
  (Unnamed, diet p_t, toStruct p_t)
  where
    p_t = patternType 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 $ prettyString t, t)
      | t <-
          Bool
            : map Signed [minBound .. maxBound]
            ++ map Unsigned [minBound .. maxBound]
            ++ map FloatType [minBound .. maxBound]
    ]

-- | The nature of something predefined.  For functions, 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] [ParamType] (RetTypeBase Size Uniqueness)
  | IntrinsicType Liftedness [TypeParamBase VName] StructType
  | IntrinsicEquality -- Special cased.

intrinsicAcc :: (VName, Intrinsic)
intrinsicAcc =
  ( acc_v,
    IntrinsicType SizeLifted [TypeParamType Unlifted t_v mempty] $
      Scalar $
        TypeVar mempty (qualName acc_v) [arg]
  )
  where
    acc_v = VName "acc" 10
    t_v = VName "t" 11
    arg = TypeArgType $ Scalar (TypeVar mempty (qualName t_v) [])

-- | If this type corresponds to the builtin "acc" type, return the
-- type of the underlying array.
isAccType :: TypeBase d u -> Maybe (TypeBase d NoUniqueness)
isAccType (Scalar (TypeVar _ (QualName [] v) [TypeArgType t]))
  | v == fst intrinsicAcc =
      Just t
isAccType _ = Nothing

-- | Find the 'VName' corresponding to a builtin.  Crashes if that
-- name cannot be found.
intrinsicVar :: Name -> VName
intrinsicVar v =
  fromMaybe bad $ find ((v ==) . baseName) $ M.keys intrinsics
  where
    bad = error $ "findBuiltin: " <> nameToString v

mkBinOp :: Name -> StructType -> Exp -> Exp -> Exp
mkBinOp op t x y =
  AppExp
    ( BinOp
        (qualName (intrinsicVar op), mempty)
        (Info t)
        (x, Info Nothing)
        (y, Info Nothing)
        mempty
    )
    (Info $ AppRes t [])

mkAdd, mkMul :: Exp -> Exp -> Exp
mkAdd = mkBinOp "+" $ Scalar $ Prim $ Signed Int64
mkMul = mkBinOp "*" $ Scalar $ Prim $ Signed Int64

-- | A map of all built-ins.
intrinsics :: M.Map VName Intrinsic
intrinsics =
  (M.fromList [intrinsicAcc] <>) $
    M.fromList $
      primOp
        ++ zipWith
          namify
          [intrinsicStart ..]
          ( [ ( "flatten",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m]
                  [Array Observe (shape [n, m]) $ t_a mempty]
                  $ RetType []
                  $ Array
                    Nonunique
                    (Shape [size n `mkMul` size m])
                    (t_a mempty)
              ),
              ( "unflatten",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m]
                  [ Scalar $ Prim $ Signed Int64,
                    Scalar $ Prim $ Signed Int64,
                    Array Observe (Shape [size n `mkMul` size m]) $ t_a mempty
                  ]
                  $ RetType []
                  $ Array Nonunique (shape [n, m]) (t_a mempty)
              ),
              ( "concat",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m]
                  [ array_a Observe $ shape [n],
                    array_a Observe $ shape [m]
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ Shape [size n `mkAdd` size m]
              ),
              ( "transpose",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m]
                  [array_a Observe $ shape [n, m]]
                  $ RetType []
                  $ array_a Nonunique
                  $ shape [m, n]
              ),
              ( "scatter",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_l]
                  [ Array Consume (shape [n]) $ t_a mempty,
                    Array Observe (shape [l]) (Prim $ Signed Int64),
                    Array Observe (shape [l]) $ t_a mempty
                  ]
                  $ RetType []
                  $ Array Unique (shape [n]) (t_a mempty)
              ),
              ( "scatter_2d",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m, sp_l]
                  [ array_a Consume $ shape [n, m],
                    Array Observe (shape [l]) (tupInt64 2),
                    Array Observe (shape [l]) $ t_a mempty
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [n, m]
              ),
              ( "scatter_3d",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m, sp_k, sp_l]
                  [ array_a Consume $ shape [n, m, k],
                    Array Observe (shape [l]) (tupInt64 3),
                    Array Observe (shape [l]) $ t_a mempty
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [n, m, k]
              ),
              ( "zip",
                IntrinsicPolyFun
                  [tp_a, tp_b, sp_n]
                  [ array_a Observe (shape [n]),
                    array_b Observe (shape [n])
                  ]
                  $ RetType []
                  $ tuple_array Unique (Scalar $ t_a mempty) (Scalar $ t_b mempty)
                  $ shape [n]
              ),
              ( "unzip",
                IntrinsicPolyFun
                  [tp_a, tp_b, sp_n]
                  [tuple_array Observe (Scalar $ t_a mempty) (Scalar $ t_b mempty) $ shape [n]]
                  $ RetType [] . Scalar . Record . M.fromList
                  $ zip tupleFieldNames [array_a Unique $ shape [n], array_b Unique $ shape [n]]
              ),
              ( "hist_1d",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m]
                  [ Scalar $ Prim $ Signed Int64,
                    array_a Consume $ shape [m],
                    Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    Array Observe (shape [n]) (tupInt64 1),
                    array_a Observe (shape [n])
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [m]
              ),
              ( "hist_2d",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m, sp_k]
                  [ Scalar $ Prim $ Signed Int64,
                    array_a Consume $ shape [m, k],
                    Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    Array Observe (shape [n]) (tupInt64 2),
                    array_a Observe (shape [n])
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [m, k]
              ),
              ( "hist_3d",
                IntrinsicPolyFun
                  [tp_a, sp_n, sp_m, sp_k, sp_l]
                  [ Scalar $ Prim $ Signed Int64,
                    array_a Consume $ shape [m, k, l],
                    Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    Array Observe (shape [n]) (tupInt64 3),
                    array_a Observe (shape [n])
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [m, k, l]
              ),
              ( "map",
                IntrinsicPolyFun
                  [tp_a, tp_b, sp_n]
                  [ Scalar (t_a mempty) `arr` Scalar (t_b Nonunique),
                    array_a Observe $ shape [n]
                  ]
                  $ RetType []
                  $ array_b Unique
                  $ shape [n]
              ),
              ( "reduce",
                IntrinsicPolyFun
                  [tp_a, sp_n]
                  [ Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    array_a Observe $ shape [n]
                  ]
                  $ RetType []
                  $ Scalar (t_a Unique)
              ),
              ( "reduce_comm",
                IntrinsicPolyFun
                  [tp_a, sp_n]
                  [ Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    array_a Observe $ shape [n]
                  ]
                  $ RetType [] (Scalar (t_a Unique))
              ),
              ( "scan",
                IntrinsicPolyFun
                  [tp_a, sp_n]
                  [ Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    array_a Observe $ shape [n]
                  ]
                  $ RetType [] (array_a Unique $ shape [n])
              ),
              ( "partition",
                IntrinsicPolyFun
                  [tp_a, sp_n]
                  [ Scalar (Prim $ Signed Int32),
                    Scalar (t_a mempty) `arr` Scalar (Prim $ Signed Int64),
                    array_a Observe $ shape [n]
                  ]
                  ( RetType [k] . Scalar $
                      tupleRecord
                        [ array_a Unique $ shape [n],
                          Array Unique (shape [k]) (Prim $ Signed Int64)
                        ]
                  )
              ),
              ( "acc_write",
                IntrinsicPolyFun
                  [sp_k, tp_a]
                  [ Scalar $ accType Consume $ array_ka mempty,
                    Scalar (Prim $ Signed Int64),
                    Scalar $ t_a Observe
                  ]
                  $ RetType []
                  $ Scalar
                  $ accType Unique (array_ka mempty)
              ),
              ( "scatter_stream",
                IntrinsicPolyFun
                  [tp_a, tp_b, sp_k, sp_n]
                  [ array_ka Consume,
                    Scalar (accType mempty (array_ka mempty))
                      `carr` ( Scalar (t_b mempty)
                                 `arr` Scalar (accType Nonunique $ array_a mempty $ shape [k])
                             ),
                    array_b Observe $ shape [n]
                  ]
                  $ RetType []
                  $ array_ka Unique
              ),
              ( "hist_stream",
                IntrinsicPolyFun
                  [tp_a, tp_b, sp_k, sp_n]
                  [ array_a Consume $ shape [k],
                    Scalar (t_a mempty) `arr` (Scalar (t_a mempty) `arr` Scalar (t_a Nonunique)),
                    Scalar $ t_a Observe,
                    Scalar (accType mempty $ array_ka mempty)
                      `carr` ( Scalar (t_b mempty)
                                 `arr` Scalar (accType Nonunique $ array_a mempty $ shape [k])
                             ),
                    array_b Observe $ shape [n]
                  ]
                  $ RetType []
                  $ array_a Unique
                  $ shape [k]
              ),
              ( "jvp2",
                IntrinsicPolyFun
                  [tp_a, tp_b]
                  [ Scalar (t_a mempty) `arr` Scalar (t_b Nonunique),
                    Scalar (t_a Observe),
                    Scalar (t_a Observe)
                  ]
                  $ RetType []
                  $ Scalar
                  $ tupleRecord [Scalar $ t_b Nonunique, Scalar $ t_b Nonunique]
              ),
              ( "vjp2",
                IntrinsicPolyFun
                  [tp_a, tp_b]
                  [ Scalar (t_a mempty) `arr` Scalar (t_b Nonunique),
                    Scalar (t_a Observe),
                    Scalar (t_b Observe)
                  ]
                  $ RetType []
                  $ Scalar
                  $ tupleRecord [Scalar $ t_b Nonunique, Scalar $ t_a Nonunique]
              )
            ]
              ++
              -- Experimental LMAD ones.
              [ ( "flat_index_2d",
                  IntrinsicPolyFun
                    [tp_a, sp_n]
                    [ array_a Observe $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64)
                    ]
                    $ RetType [m, k]
                    $ array_a Nonunique
                    $ shape [m, k]
                ),
                ( "flat_update_2d",
                  IntrinsicPolyFun
                    [tp_a, sp_n, sp_k, sp_l]
                    [ array_a Consume $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      array_a Observe $ shape [k, l]
                    ]
                    $ RetType []
                    $ array_a Unique
                    $ shape [n]
                ),
                ( "flat_index_3d",
                  IntrinsicPolyFun
                    [tp_a, sp_n]
                    [ array_a Observe $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64)
                    ]
                    $ RetType [m, k, l]
                    $ array_a Nonunique
                    $ shape [m, k, l]
                ),
                ( "flat_update_3d",
                  IntrinsicPolyFun
                    [tp_a, sp_n, sp_k, sp_l, sp_p]
                    [ array_a Consume $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      array_a Observe $ shape [k, l, p]
                    ]
                    $ RetType []
                    $ array_a Unique
                    $ shape [n]
                ),
                ( "flat_index_4d",
                  IntrinsicPolyFun
                    [tp_a, sp_n]
                    [ array_a Observe $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64)
                    ]
                    $ RetType [m, k, l, p]
                    $ array_a Nonunique
                    $ shape [m, k, l, p]
                ),
                ( "flat_update_4d",
                  IntrinsicPolyFun
                    [tp_a, sp_n, sp_k, sp_l, sp_p, sp_q]
                    [ array_a Consume $ shape [n],
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      Scalar (Prim $ Signed Int64),
                      array_a Observe $ shape [k, l, p, q]
                    ]
                    $ RetType []
                    $ array_a Unique
                    $ shape [n]
                )
              ]
          )
  where
    primOp =
      zipWith namify [20 ..] $
        map primFun (M.toList Primitive.primFuns)
          ++ 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
            intrinsicPrim
            ( map Signed [minBound .. maxBound]
                ++ map Unsigned [minBound .. maxBound]
                ++ map FloatType [minBound .. maxBound]
                ++ [Bool]
            )
          ++
          -- This overrides the ! from Primitive.

          -- 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]

    intrinsicStart = 1 + baseTag (fst $ last primOp)

    [a, b, n, m, k, l, p, q] = zipWith VName (map nameFromString ["a", "b", "n", "m", "k", "l", "p", "q"]) [0 ..]

    t_a u = TypeVar u (qualName a) []
    array_a u s = Array u s $ t_a mempty
    tp_a = TypeParamType Unlifted a mempty

    t_b u = TypeVar u (qualName b) []
    array_b u s = Array u s $ t_b mempty
    tp_b = TypeParamType Unlifted b mempty

    [sp_n, sp_m, sp_k, sp_l, sp_p, sp_q] = map (`TypeParamDim` mempty) [n, m, k, l, p, q]

    size = flip sizeFromName mempty . qualName
    shape = Shape . map size

    tuple_array u x y s =
      Array u s (Record (M.fromList $ zip tupleFieldNames [x, y]))

    arr x y = Scalar $ Arrow mempty Unnamed Observe x (RetType [] y)
    carr x y = Scalar $ Arrow mempty Unnamed Consume x (RetType [] y)

    array_ka u = Array u (Shape [sizeFromName (qualName k) mempty]) $ t_a mempty

    accType u t =
      TypeVar u (qualName (fst intrinsicAcc)) [TypeArgType t]

    namify i (x, y) = (VName (nameFromString x) i, y)

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

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

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

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

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

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

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

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

    intrinsicPrim t = (prettyString t, IntrinsicType Unlifted [] $ Scalar $ Prim 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
      pure (prettyString 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 1 =
      Prim $ Signed Int64
    tupInt64 x =
      tupleRecord $ replicate x $ Scalar $ Prim $ Signed Int64

-- | Is this include part of the built-in prelude?
isBuiltin :: FilePath -> Bool
isBuiltin = (== "/prelude") . takeDirectory

-- | Is the position of this thing builtin as per 'isBuiltin'?  Things
-- without location are considered not built-in.
isBuiltinLoc :: (Located a) => a -> Bool
isBuiltinLoc x =
  case locOf x of
    NoLoc -> False
    Loc pos _ -> isBuiltin $ posFile pos

-- | 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

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

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

modExpImports :: ModExpBase f vn -> [(String, Loc)]
modExpImports ModVar {} = []
modExpImports (ModParens p _) = modExpImports p
modExpImports (ModImport f _ loc) = [(f, locOf 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 :: ProgBase Info VName -> S.Set VName
progModuleTypes prog = foldMap reach mtypes_used
  where
    -- Fixed point iteration.
    reach v = S.singleton v <> maybe mempty (foldMap reach) (M.lookup v reachable_from_mtype)

    reachable_from_mtype = foldMap onDec $ progDecs prog
      where
        onDec OpenDec {} = mempty
        onDec ModDec {} = mempty
        onDec (ModTypeDec sb) =
          M.singleton (modTypeName sb) (onModTypeExp (modTypeExp sb))
        onDec TypeDec {} = mempty
        onDec ValDec {} = mempty
        onDec (LocalDec d _) = onDec d
        onDec ImportDec {} = mempty

        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v
        onModTypeExp (ModTypeParens e _) = onModTypeExp e
        onModTypeExp (ModTypeSpecs ss _) = foldMap onSpec ss
        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e
        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp e2

        onSpec ValSpec {} = mempty
        onSpec TypeSpec {} = mempty
        onSpec TypeAbbrSpec {} = mempty
        onSpec (ModSpec vn e _ _) = S.singleton vn <> onModTypeExp e
        onSpec (IncludeSpec e _) = onModTypeExp e

    mtypes_used = foldMap onDec $ progDecs prog
      where
        onDec (OpenDec x _) = onModExp x
        onDec (ModDec md) =
          maybe mempty (onModTypeExp . fst) (modType md) <> onModExp (modExp md)
        onDec ModTypeDec {} = 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 <> onModTypeExp se
        onModExp (ModLambda p r me _) =
          onModParam p <> maybe mempty (onModTypeExp . fst) r <> onModExp me

        onModParam = onModTypeExp . modParamType

        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v
        onModTypeExp (ModTypeParens e _) = onModTypeExp e
        onModTypeExp ModTypeSpecs {} = mempty
        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e
        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp 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 prettyString operators) operators
  where
    s' = nameToString s
    operators :: [BinOp]
    operators = [minBound .. maxBound :: BinOp]

-- | Find instances of typed holes in the program.
progHoles :: ProgBase Info VName -> [(Loc, StructType)]
progHoles = foldMap holesInDec . progDecs
  where
    holesInDec (ValDec vb) = holesInExp $ valBindBody vb
    holesInDec (ModDec me) = holesInModExp $ modExp me
    holesInDec (OpenDec me _) = holesInModExp me
    holesInDec (LocalDec d _) = holesInDec d
    holesInDec TypeDec {} = mempty
    holesInDec ModTypeDec {} = mempty
    holesInDec ImportDec {} = mempty

    holesInModExp (ModDecs ds _) = foldMap holesInDec ds
    holesInModExp (ModParens me _) = holesInModExp me
    holesInModExp (ModApply x y _ _ _) = holesInModExp x <> holesInModExp y
    holesInModExp (ModAscript me _ _ _) = holesInModExp me
    holesInModExp (ModLambda _ _ me _) = holesInModExp me
    holesInModExp ModVar {} = mempty
    holesInModExp ModImport {} = mempty

    holesInExp = flip execState mempty . onExp

    onExp e@(Hole (Info t) loc) = do
      modify ((locOf loc, toStruct t) :)
      pure e
    onExp e = astMap (identityMapper {mapOnExp = onExp}) e

-- | Strip semantically irrelevant stuff from the top level of
-- expression.  This is used to provide a slightly fuzzy notion of
-- expression equality.
--
-- Ideally we'd implement unification on a simpler representation that
-- simply didn't allow us.
stripExp :: Exp -> Maybe Exp
stripExp (Parens e _) = stripExp e `mplus` Just e
stripExp (Assert _ e _ _) = stripExp e `mplus` Just e
stripExp (Attr _ e _) = stripExp e `mplus` Just e
stripExp (Ascript e _ _) = stripExp e `mplus` Just e
stripExp _ = Nothing

similarSlices :: Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices slice1 slice2
  | length slice1 == length slice2 = do
      concat <$> zipWithM match slice1 slice2
  | otherwise = Nothing
  where
    match (DimFix e1) (DimFix e2) = Just [(e1, e2)]
    match (DimSlice a1 b1 c1) (DimSlice a2 b2 c2) =
      concat <$> sequence [pair (a1, a2), pair (b1, b2), pair (c1, c2)]
    match _ _ = Nothing
    pair (Nothing, Nothing) = Just []
    pair (Just x, Just y) = Just [(x, y)]
    pair _ = Nothing

-- | If these two expressions are structurally similar at top level as
-- sizes, produce their subexpressions (which are not necessarily
-- similar, but you can check for that!).  This is the machinery
-- underlying expresssion unification.
similarExps :: Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps e1 e2 | bareExp e1 == bareExp e2 = Just []
similarExps e1 e2 | Just e1' <- stripExp e1 = similarExps e1' e2
similarExps e1 e2 | Just e2' <- stripExp e2 = similarExps e1 e2'
similarExps
  (AppExp (BinOp (op1, _) _ (x1, _) (y1, _) _) _)
  (AppExp (BinOp (op2, _) _ (x2, _) (y2, _) _) _)
    | op1 == op2 = Just [(x1, x2), (y1, y2)]
similarExps (AppExp (Apply f1 args1 _) _) (AppExp (Apply f2 args2 _) _)
  | f1 == f2 = Just $ zip (map snd $ NE.toList args1) (map snd $ NE.toList args2)
similarExps (AppExp (Index arr1 slice1 _) _) (AppExp (Index arr2 slice2 _) _)
  | arr1 == arr2,
    length slice1 == length slice2 =
      similarSlices slice1 slice2
similarExps (TupLit es1 _) (TupLit es2 _)
  | length es1 == length es2 =
      Just $ zip es1 es2
similarExps (RecordLit fs1 _) (RecordLit fs2 _)
  | length fs1 == length fs2 =
      zipWithM onFields fs1 fs2
  where
    onFields (RecordFieldExplicit n1 fe1 _) (RecordFieldExplicit n2 fe2 _)
      | n1 == n2 = Just (fe1, fe2)
    onFields (RecordFieldImplicit vn1 ty1 _) (RecordFieldImplicit vn2 ty2 _) =
      Just (Var (qualName vn1) ty1 mempty, Var (qualName vn2) ty2 mempty)
    onFields _ _ = Nothing
similarExps (ArrayLit es1 _ _) (ArrayLit es2 _ _)
  | length es1 == length es2 =
      Just $ zip es1 es2
similarExps (Project field1 e1 _ _) (Project field2 e2 _ _)
  | field1 == field2 =
      Just [(e1, e2)]
similarExps (Negate e1 _) (Negate e2 _) =
  Just [(e1, e2)]
similarExps (Not e1 _) (Not e2 _) =
  Just [(e1, e2)]
similarExps (Constr n1 es1 _ _) (Constr n2 es2 _ _)
  | length es1 == length es2,
    n1 == n2 =
      Just $ zip es1 es2
similarExps (Update e1 slice1 e'1 _) (Update e2 slice2 e'2 _) =
  ([(e1, e2), (e'1, e'2)] ++) <$> similarSlices slice1 slice2
similarExps (RecordUpdate e1 names1 e'1 _ _) (RecordUpdate e2 names2 e'2 _ _)
  | names1 == names2 =
      Just [(e1, e2), (e'1, e'2)]
similarExps (OpSection op1 _ _) (OpSection op2 _ _)
  | op1 == op2 = Just []
similarExps (OpSectionLeft op1 _ x1 _ _ _) (OpSectionLeft op2 _ x2 _ _ _)
  | op1 == op2 = Just [(x1, x2)]
similarExps (OpSectionRight op1 _ x1 _ _ _) (OpSectionRight op2 _ x2 _ _ _)
  | op1 == op2 = Just [(x1, x2)]
similarExps (ProjectSection names1 _ _) (ProjectSection names2 _ _)
  | names1 == names2 = Just []
similarExps (IndexSection slice1 _ _) (IndexSection slice2 _ _) =
  similarSlices slice1 slice2
similarExps _ _ = Nothing

-- | An identifier with type- and aliasing information.
type Ident = IdentBase Info VName

-- | An index with type information.
type DimIndex = DimIndexBase Info VName

-- | A slice with type information.
type Slice = SliceBase Info VName

-- | An expression with type information.
type Exp = ExpBase Info VName

-- | An application expression with type information.
type AppExp = AppExpBase Info VName

-- | A pattern with type information.
type Pat = PatBase Info VName

-- | An constant declaration with type information.
type ValBind = ValBindBase Info VName

-- | A type binding with type information.
type TypeBind = TypeBindBase Info VName

-- | A type-checked module binding.
type ModBind = ModBindBase Info VName

-- | A type-checked module type binding.
type ModTypeBind = ModTypeBindBase Info VName

-- | A type-checked module expression.
type ModExp = ModExpBase Info VName

-- | A type-checked module parameter.
type ModParam = ModParamBase Info VName

-- | A type-checked module type expression.
type ModTypeExp = ModTypeExpBase Info VName

-- | A type-checked declaration.
type Dec = DecBase Info VName

-- | A type-checked specification.
type Spec = SpecBase Info VName

-- | An Futhark program with type information.
type Prog = ProgBase Info VName

-- | A known type arg with shape annotations.
type StructTypeArg = TypeArg Size

-- | A type-checked type parameter.
type TypeParam = TypeParamBase VName

-- | A known scalar type with no shape annotations.
type ScalarType = ScalarTypeBase ()

-- | A type-checked case (of a match expression).
type Case = CaseBase Info VName

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

-- | An unchecked type expression.
type UncheckedTypeExp = TypeExp NoInfo Name

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

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

-- | A slice with no type annotations.
type UncheckedSlice = SliceBase 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 UncheckedModTypeExp = ModTypeExpBase NoInfo Name

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

-- | A pattern with no type annotations.
type UncheckedPat = PatBase NoInfo Name

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

-- | A type binding with no type annotations.
type UncheckedTypeBind = TypeBindBase NoInfo Name

-- | A module type binding with no type annotations.
type UncheckedModTypeBind = ModTypeBindBase NoInfo Name

-- | A module binding with no type annotations.
type UncheckedModBind = ModBindBase 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