content: -- | Type checking of patterns.
  module Language.Futhark.TypeChecker.Terms.Pat
    ( binding,
      bindingParams,
      bindingPat,
      bindingIdent,
      bindingSizes,
    )
  where
  
  import Control.Monad
  import Data.Bifunctor
  import Data.Either
  import Data.List (find, isPrefixOf, sort, sortBy)
  import Data.Map.Strict qualified as M
  import Data.Maybe
  import Data.Ord (comparing)
  import Data.Set qualified as S
  import Futhark.Util.Pretty hiding (group, space)
  import Language.Futhark
  import Language.Futhark.TypeChecker.Monad hiding (BoundV)
  import Language.Futhark.TypeChecker.Terms.Monad
  import Language.Futhark.TypeChecker.Types
  import Language.Futhark.TypeChecker.Unify hiding (Usage)
  import Prelude hiding (mod)
  
  nonrigidFor :: [(SizeBinder VName, QualName VName)] -> StructType -> StructType
  nonrigidFor [] = id -- Minor optimisation.
  nonrigidFor sizes = applySubst onDim
    where
      onDim v
        | Just (s, v') <- find ((== v) . sizeName . fst) sizes =
            Just $ ExpSubst $ sizeFromName v' (srclocOf s)
      onDim _ = Nothing
  
  -- | Bind these identifiers locally while running the provided action.
  binding ::
    [Ident StructType] ->
    TermTypeM a ->
    TermTypeM a
  binding idents m =
    localScope (`bindVars` idents) $ do
      -- Those identifiers that can potentially also be sizes are
      -- added as type constraints.  This is necessary so that we
      -- can properly detect scope violations during unification.
      -- We do this for *all* identifiers, not just those that are
      -- integers, because they may become integers later due to
      -- inference...
      forM_ idents $ \ident ->
        constrain (identName ident) $ ParamSize $ locOf ident
      m
    where
      bindVars = foldl bindVar
  
      bindVar scope (Ident name (Info tp) _) =
        scope
          { scopeVtable =
              M.insert name (BoundV [] tp) $ scopeVtable scope
          }
  
  bindingTypes ::
    [Either (VName, TypeBinding) (VName, Constraint)] ->
    TermTypeM a ->
    TermTypeM a
  bindingTypes types m = do
    lvl <- curLevel
    modifyConstraints (<> M.map (lvl,) (M.fromList constraints))
    localScope extend m
    where
      (tbinds, constraints) = partitionEithers types
      extend scope =
        scope
          { scopeTypeTable = M.fromList tbinds <> scopeTypeTable scope
          }
  
  bindingTypeParams :: [TypeParam] -> TermTypeM a -> TermTypeM a
  bindingTypeParams tparams =
    binding (mapMaybe typeParamIdent tparams)
      . bindingTypes (concatMap typeParamType tparams)
    where
      typeParamType (TypeParamType l v loc) =
        [ Left (v, TypeAbbr l [] $ RetType [] $ Scalar (TypeVar mempty (qualName v) [])),
          Right (v, ParamType l $ locOf loc)
        ]
      typeParamType (TypeParamDim v loc) =
        [Right (v, ParamSize $ locOf loc)]
  
  typeParamIdent :: TypeParam -> Maybe (Ident StructType)
  typeParamIdent (TypeParamDim v loc) =
    Just $ Ident v (Info $ Scalar $ Prim $ Signed Int64) loc
  typeParamIdent _ = Nothing
  
  -- | Bind @let@-bound sizes.  This is usually followed by 'bindingPat'
  -- immediately afterwards.
  bindingSizes :: [SizeBinder VName] -> TermTypeM a -> TermTypeM a
  bindingSizes [] m = m -- Minor optimisation.
  bindingSizes sizes m = binding (map sizeWithType sizes) m
    where
      sizeWithType size =
        Ident (sizeName size) (Info (Scalar (Prim (Signed Int64)))) (srclocOf size)
  
  -- | Bind a single term-level identifier.
  bindingIdent ::
    IdentBase NoInfo VName StructType ->
    StructType ->
    (Ident StructType -> TermTypeM a) ->
    TermTypeM a
  bindingIdent (Ident v NoInfo vloc) t m = do
    let ident = Ident v (Info t) vloc
    binding [ident] $ m ident
  
  -- All this complexity is just so we can handle un-suffixed numeric
  -- literals in patterns.
  patLitMkType :: PatLit -> SrcLoc -> TermTypeM ParamType
  patLitMkType (PatLitInt _) loc = do
    t <- newTypeVar loc "t"
    mustBeOneOf anyNumberType (mkUsage loc "integer literal") (toStruct t)
    pure t
  patLitMkType (PatLitFloat _) loc = do
    t <- newTypeVar loc "t"
    mustBeOneOf anyFloatType (mkUsage loc "float literal") (toStruct t)
    pure t
  patLitMkType (PatLitPrim v) _ =
    pure $ Scalar $ Prim $ primValueType v
  
  checkPat' ::
    [(SizeBinder VName, QualName VName)] ->
    PatBase NoInfo VName ParamType ->
    Inferred ParamType ->
    TermTypeM (Pat ParamType)
  checkPat' sizes (PatParens p loc) t =
    PatParens <$> checkPat' sizes p t <*> pure loc
  checkPat' sizes (PatAttr attr p loc) t =
    PatAttr <$> checkAttr attr <*> checkPat' sizes p t <*> pure loc
  checkPat' _ (Id name NoInfo loc) (Ascribed t) =
    pure $ Id name (Info t) loc
  checkPat' _ (Id name NoInfo loc) NoneInferred = do
    t <- newTypeVar loc "t"
    pure $ Id name (Info t) loc
  checkPat' _ (Wildcard _ loc) (Ascribed t) =
    pure $ Wildcard (Info t) loc
  checkPat' _ (Wildcard NoInfo loc) NoneInferred = do
    t <- newTypeVar loc "t"
    pure $ Wildcard (Info t) loc
  checkPat' sizes p@(TuplePat ps loc) (Ascribed t)
    | Just ts <- isTupleRecord t,
      length ts == length ps =
        TuplePat
          <$> zipWithM (checkPat' sizes) ps (map Ascribed ts)
          <*> pure loc
    | otherwise = do
        ps_t <- replicateM (length ps) (newTypeVar loc "t")
        unify (mkUsage loc "matching a tuple pattern") (Scalar (tupleRecord ps_t)) (toStruct t)
        checkPat' sizes p $ Ascribed $ toParam Observe $ Scalar $ tupleRecord ps_t
  checkPat' sizes (TuplePat ps loc) NoneInferred =
    TuplePat <$> mapM (\p -> checkPat' sizes p NoneInferred) ps <*> pure loc
  checkPat' _ (RecordPat p_fs _) _
    | Just (L loc f, _) <- find (("_" `isPrefixOf`) . nameToString . unLoc . fst) p_fs =
        typeError loc mempty $
          "Underscore-prefixed fields are not allowed."
            </> "Did you mean"
            <> dquotes (pretty (drop 1 (nameToString f)) <> "=_")
            <> "?"
  checkPat' sizes p@(RecordPat p_fs loc) (Ascribed t)
    | Scalar (Record t_fs) <- t,
      p_fs' <- sortBy (comparing fst) p_fs,
      t_fs' <- sortBy (comparing fst) (M.toList t_fs),
      map fst t_fs' == map (unLoc . fst) p_fs' =
        RecordPat <$> zipWithM check p_fs' t_fs' <*> pure loc
    | otherwise = do
        p_fs' <- traverse (const $ newTypeVar loc "t") $ M.fromList $ map (first unLoc) p_fs
  
        when (sort (M.keys p_fs') /= sort (map (unLoc . fst) p_fs)) $
          typeError loc mempty $
            "Duplicate fields in record pattern" <+> pretty p <> "."
  
        unify (mkUsage loc "matching a record pattern") (Scalar (Record p_fs')) (toStruct t)
        checkPat' sizes p $ Ascribed $ toParam Observe $ Scalar (Record p_fs')
    where
      check (L f_loc f, p_f) (_, t_f) = (L f_loc f,) <$> checkPat' sizes p_f (Ascribed t_f)
  checkPat' sizes (RecordPat fs loc) NoneInferred =
    RecordPat . M.toList
      <$> traverse (\p -> checkPat' sizes p NoneInferred) (M.fromList fs)
      <*> pure loc
  checkPat' sizes (PatAscription p t loc) maybe_outer_t = do
    (t', st, _) <- checkTypeExpNonrigid t
  
    case maybe_outer_t of
      Ascribed outer_t -> do
        let st_forunify = nonrigidFor sizes $ toStruct st
        unify (mkUsage loc "explicit type ascription") st_forunify (toStruct outer_t)
  
        PatAscription
          <$> checkPat' sizes p (Ascribed (resToParam st))
          <*> pure t'
          <*> pure loc
      NoneInferred ->
        PatAscription
          <$> checkPat' sizes p (Ascribed (resToParam st))
          <*> pure t'
          <*> pure loc
  checkPat' _ (PatLit l NoInfo loc) (Ascribed t) = do
    t' <- patLitMkType l loc
    unify (mkUsage loc "matching against literal") (toStruct t') (toStruct t)
    pure $ PatLit l (Info t') loc
  checkPat' _ (PatLit l NoInfo loc) NoneInferred = do
    t' <- patLitMkType l loc
    pure $ PatLit l (Info t') loc
  checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs)))
    | Just ts <- M.lookup n cs = do
        when (length ps /= length ts) $
          typeError loc mempty $
            "Pattern #"
              <> pretty n
              <> " expects"
                <+> pretty (length ps)
                <+> "constructor arguments, but type provides"
                <+> pretty (length ts)
                <+> "arguments."
        ps' <- zipWithM (checkPat' sizes) ps $ map Ascribed ts
        pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc
  checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed t) = do
    t' <- newTypeVar loc "t"
    ps' <- forM ps $ \p -> do
      p_t <- newTypeVar (srclocOf p) "t"
      checkPat' sizes p $ Ascribed p_t
    mustHaveConstr usage n (toStruct t') (patternStructType <$> ps')
    unify usage t' (toStruct t)
    pure $ PatConstr n (Info t) ps' loc
    where
      usage = mkUsage loc "matching against constructor"
  checkPat' sizes (PatConstr n NoInfo ps loc) NoneInferred = do
    ps' <- mapM (\p -> checkPat' sizes p NoneInferred) ps
    t <- newTypeVar loc "t"
    mustHaveConstr usage n (toStruct t) (patternStructType <$> ps')
    pure $ PatConstr n (Info t) ps' loc
    where
      usage = mkUsage loc "matching against constructor"
  
  checkPat ::
    [(SizeBinder VName, QualName VName)] ->
    PatBase NoInfo VName (TypeBase Size u) ->
    Inferred StructType ->
    (Pat ParamType -> TermTypeM a) ->
    TermTypeM a
  checkPat sizes p t m = do
    p' <-
      onFailure (CheckingPat (fmap toStruct p) t) $
        checkPat' sizes (fmap (toParam Observe) p) (fmap (toParam Observe) t)
  
    let explicit = mustBeExplicitInType $ patternStructType p'
  
    case filter ((`S.member` explicit) . sizeName . fst) sizes of
      (size, _) : _ ->
        typeError size mempty $
          "Cannot bind"
            <+> pretty size
            <+> "as it is never used as the size of a concrete (non-function) value."
      [] ->
        m p'
  
  -- | Check and bind a @let@-pattern.
  bindingPat ::
    [SizeBinder VName] ->
    PatBase NoInfo VName (TypeBase Size u) ->
    StructType ->
    (Pat ParamType -> TermTypeM a) ->
    TermTypeM a
  bindingPat sizes p t m = do
    substs <- mapM mkSizeSubst sizes
    checkPat substs p (Ascribed t) $ \p' -> binding (patIdents (fmap toStruct p')) $
      case filter ((`S.notMember` fvVars (freeInPat p')) . sizeName) sizes of
        [] -> m p'
        size : _ -> unusedSize size
    where
      mkSizeSubst v = do
        v' <- newID $ baseName $ sizeName v
        constrain v' . Size Nothing $
          mkUsage v "ambiguous size of bound expression"
        pure (v, qualName v')
  
  -- | Check and bind type and value parameters.
  bindingParams ::
    [TypeParam] ->
    [PatBase NoInfo VName ParamType] ->
    ([Pat ParamType] -> TermTypeM a) ->
    TermTypeM a
  bindingParams tps orig_ps m = bindingTypeParams tps $ do
    let descend ps' (p : ps) =
          checkPat [] p NoneInferred $ \p' ->
            binding (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps
        descend ps' [] = m $ reverse ps'
  
    incLevel $ descend [] orig_ps
  