{-# LANGUAGE Strict #-}

-- | Internalising bindings.
module Futhark.Internalise.Bindings
  ( internaliseAttrs,
    internaliseAttr,
    bindingFParams,
    bindingLoopParams,
    bindingLambdaParams,
    stmPat,
  )
where

import Control.Monad
import Control.Monad.Free (Free (..))
import Control.Monad.Reader
import Data.Bifunctor
import Data.Foldable (toList)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.IR.SOACS qualified as I
import Futhark.Internalise.Monad
import Futhark.Internalise.TypesValues
import Futhark.Util
import Language.Futhark as E hiding (matchDims)

internaliseAttr :: E.AttrInfo VName -> InternaliseM I.Attr
internaliseAttr :: AttrInfo VName -> InternaliseM Attr
internaliseAttr (E.AttrAtom (E.AtomName Name
v) SrcLoc
_) =
  Attr -> InternaliseM Attr
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attr -> InternaliseM Attr) -> Attr -> InternaliseM Attr
forall a b. (a -> b) -> a -> b
$ Name -> Attr
I.AttrName Name
v
internaliseAttr (E.AttrAtom (E.AtomInt Integer
x) SrcLoc
_) =
  Attr -> InternaliseM Attr
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attr -> InternaliseM Attr) -> Attr -> InternaliseM Attr
forall a b. (a -> b) -> a -> b
$ Integer -> Attr
I.AttrInt Integer
x
internaliseAttr (E.AttrComp Name
f [AttrInfo VName]
attrs SrcLoc
_) =
  Name -> [Attr] -> Attr
I.AttrComp Name
f ([Attr] -> Attr) -> InternaliseM [Attr] -> InternaliseM Attr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (AttrInfo VName -> InternaliseM Attr)
-> [AttrInfo VName] -> InternaliseM [Attr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM AttrInfo VName -> InternaliseM Attr
internaliseAttr [AttrInfo VName]
attrs

internaliseAttrs :: [E.AttrInfo VName] -> InternaliseM I.Attrs
internaliseAttrs :: [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs = ([Attr] -> Attrs) -> InternaliseM [Attr] -> InternaliseM Attrs
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Attrs] -> Attrs
forall a. Monoid a => [a] -> a
mconcat ([Attrs] -> Attrs) -> ([Attr] -> [Attrs]) -> [Attr] -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Attr -> Attrs) -> [Attr] -> [Attrs]
forall a b. (a -> b) -> [a] -> [b]
map Attr -> Attrs
I.oneAttr) (InternaliseM [Attr] -> InternaliseM Attrs)
-> ([AttrInfo VName] -> InternaliseM [Attr])
-> [AttrInfo VName]
-> InternaliseM Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AttrInfo VName -> InternaliseM Attr)
-> [AttrInfo VName] -> InternaliseM [Attr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM AttrInfo VName -> InternaliseM Attr
internaliseAttr

treeLike :: Tree a -> [b] -> Tree b
treeLike :: forall a b. Tree a -> [b] -> Tree b
treeLike (Pure a
_) [b
b] = b -> Free [] b
forall (f :: * -> *) a. a -> Free f a
Pure b
b
treeLike (Pure a
_) [b]
_ = [Char] -> Free [] b
forall a. HasCallStack => [Char] -> a
error [Char]
"treeLike: invalid input"
treeLike (Free [Free [] a]
ls) [b]
bs = [Free [] b] -> Free [] b
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free ([Free [] b] -> Free [] b) -> [Free [] b] -> Free [] b
forall a b. (a -> b) -> a -> b
$ (Free [] a -> [b] -> Free [] b)
-> [Free [] a] -> [[b]] -> [Free [] b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Free [] a -> [b] -> Free [] b
forall a b. Tree a -> [b] -> Tree b
treeLike [Free [] a]
ls ([Int] -> [b] -> [[b]]
forall a. [Int] -> [a] -> [[a]]
chunks ((Free [] a -> Int) -> [Free [] a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Free [] a -> Int
forall a. Free [] a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Free [] a]
ls) [b]
bs)

bindingFParams ::
  [E.TypeParam] ->
  [E.Pat E.ParamType] ->
  ([I.FParam I.SOACS] -> [[Tree (I.FParam I.SOACS)]] -> InternaliseM a) ->
  InternaliseM a
bindingFParams :: forall a.
[TypeParam]
-> [Pat ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParam]
tparams [Pat ParamType]
params [FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a
m = do
  [[(Ident ParamType, [AttrInfo VName])]]
flattened_params <- (Pat ParamType
 -> InternaliseM [(Ident ParamType, [AttrInfo VName])])
-> [Pat ParamType]
-> InternaliseM [[(Ident ParamType, [AttrInfo VName])]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> InternaliseM [(Ident ParamType, [AttrInfo VName])]
forall (m :: * -> *) u.
MonadFreshNames m =>
Pat (TypeBase Size u)
-> m [(Ident (TypeBase Size u), [AttrInfo VName])]
flattenPat [Pat ParamType]
params
  let params_idents :: [(Ident ParamType, [AttrInfo VName])]
params_idents = [[(Ident ParamType, [AttrInfo VName])]]
-> [(Ident ParamType, [AttrInfo VName])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Ident ParamType, [AttrInfo VName])]]
flattened_params
  [[Tree (TypeBase Shape Uniqueness)]]
params_ts <-
    [ParamType] -> InternaliseM [[Tree (TypeBase Shape Uniqueness)]]
internaliseParamTypes ([ParamType] -> InternaliseM [[Tree (TypeBase Shape Uniqueness)]])
-> [ParamType] -> InternaliseM [[Tree (TypeBase Shape Uniqueness)]]
forall a b. (a -> b) -> a -> b
$
      ((Ident ParamType, [AttrInfo VName]) -> ParamType)
-> [(Ident ParamType, [AttrInfo VName])] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (Info ParamType -> ParamType
forall a. Info a -> a
E.unInfo (Info ParamType -> ParamType)
-> ((Ident ParamType, [AttrInfo VName]) -> Info ParamType)
-> (Ident ParamType, [AttrInfo VName])
-> ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident ParamType -> Info ParamType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
E.identType (Ident ParamType -> Info ParamType)
-> ((Ident ParamType, [AttrInfo VName]) -> Ident ParamType)
-> (Ident ParamType, [AttrInfo VName])
-> Info ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ident ParamType, [AttrInfo VName]) -> Ident ParamType
forall a b. (a, b) -> a
fst) [(Ident ParamType, [AttrInfo VName])]
params_idents
  let num_param_idents :: [Int]
num_param_idents = ([(Ident ParamType, [AttrInfo VName])] -> Int)
-> [[(Ident ParamType, [AttrInfo VName])]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [(Ident ParamType, [AttrInfo VName])] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[(Ident ParamType, [AttrInfo VName])]]
flattened_params

  let shape_params :: [Param (TypeBase shape u)]
shape_params = [Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
forall a. Monoid a => a
mempty VName
v (TypeBase shape u -> Param (TypeBase shape u))
-> TypeBase shape u -> Param (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.int64 | E.TypeParamDim VName
v SrcLoc
_ <- [TypeParam]
tparams]
      shape_subst :: Map VName [SubExp]
shape_subst = [(VName, [SubExp])] -> Map VName [SubExp]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Param (TypeBase Any Any) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Any Any)
p, [VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Any Any) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Any Any)
p]) | Param (TypeBase Any Any)
p <- [Param (TypeBase Any Any)]
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params]
  [(Ident ParamType, [AttrInfo VName])]
-> [TypeBase Shape Uniqueness]
-> ([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
-> InternaliseM a
forall t a.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> ([Params t] -> InternaliseM a) -> InternaliseM a
bindingFlatPat [(Ident ParamType, [AttrInfo VName])]
params_idents (([Tree (TypeBase Shape Uniqueness)] -> [TypeBase Shape Uniqueness])
-> [[Tree (TypeBase Shape Uniqueness)]]
-> [TypeBase Shape Uniqueness]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Tree (TypeBase Shape Uniqueness) -> [TypeBase Shape Uniqueness])
-> [Tree (TypeBase Shape Uniqueness)]
-> [TypeBase Shape Uniqueness]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree (TypeBase Shape Uniqueness) -> [TypeBase Shape Uniqueness]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (TypeBase Shape Uniqueness)]]
params_ts) (([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
 -> InternaliseM a)
-> ([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
-> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \[Params (TypeBase Shape Uniqueness)]
valueparams -> do
    let ([Param (TypeBase shape u)]
certparams, [Params (TypeBase Shape Uniqueness)]
valueparams') =
          ([[Param (TypeBase shape u)]] -> [Param (TypeBase shape u)])
-> ([[Param (TypeBase shape u)]],
    [Params (TypeBase Shape Uniqueness)])
-> ([Param (TypeBase shape u)],
    [Params (TypeBase Shape Uniqueness)])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [[Param (TypeBase shape u)]] -> [Param (TypeBase shape u)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (([[Param (TypeBase shape u)]],
  [Params (TypeBase Shape Uniqueness)])
 -> ([Param (TypeBase shape u)],
     [Params (TypeBase Shape Uniqueness)]))
-> ([[Param (TypeBase shape u)]],
    [Params (TypeBase Shape Uniqueness)])
-> ([Param (TypeBase shape u)],
    [Params (TypeBase Shape Uniqueness)])
forall a b. (a -> b) -> a -> b
$ [([Param (TypeBase shape u)], Params (TypeBase Shape Uniqueness))]
-> ([[Param (TypeBase shape u)]],
    [Params (TypeBase Shape Uniqueness)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([Param (TypeBase shape u)], Params (TypeBase Shape Uniqueness))]
 -> ([[Param (TypeBase shape u)]],
     [Params (TypeBase Shape Uniqueness)]))
-> [([Param (TypeBase shape u)],
     Params (TypeBase Shape Uniqueness))]
-> ([[Param (TypeBase shape u)]],
    [Params (TypeBase Shape Uniqueness)])
forall a b. (a -> b) -> a -> b
$ (Params (TypeBase Shape Uniqueness)
 -> ([Param (TypeBase shape u)],
     Params (TypeBase Shape Uniqueness)))
-> [Params (TypeBase Shape Uniqueness)]
-> [([Param (TypeBase shape u)],
     Params (TypeBase Shape Uniqueness))]
forall a b. (a -> b) -> [a] -> [b]
map Params (TypeBase Shape Uniqueness)
-> ([Param (TypeBase shape u)], Params (TypeBase Shape Uniqueness))
forall {shape} {u} {shape} {u}.
[Param (TypeBase shape u)]
-> ([Param (TypeBase shape u)], [Param (TypeBase shape u)])
fixAccParams [Params (TypeBase Shape Uniqueness)]
valueparams
        all_params :: Params (TypeBase Shape Uniqueness)
all_params = Params (TypeBase Shape Uniqueness)
forall {shape} {u}. [Param (TypeBase shape u)]
certparams Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
forall a. [a] -> [a] -> [a]
++ Params (TypeBase Shape Uniqueness)
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
forall a. [a] -> [a] -> [a]
++ [Params (TypeBase Shape Uniqueness)]
-> Params (TypeBase Shape Uniqueness)
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Params (TypeBase Shape Uniqueness)]
valueparams'
    Scope SOACS -> InternaliseM a -> InternaliseM a
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
I.localScope (Params (TypeBase Shape Uniqueness) -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
I.scopeOfFParams Params (TypeBase Shape Uniqueness)
all_params) (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
      Map VName [SubExp] -> InternaliseM a -> InternaliseM a
forall a. Map VName [SubExp] -> InternaliseM a -> InternaliseM a
substitutingVars Map VName [SubExp]
shape_subst (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ do
        let values_grouped_by_params :: [[Params (TypeBase Shape Uniqueness)]]
values_grouped_by_params = [Int]
-> [Params (TypeBase Shape Uniqueness)]
-> [[Params (TypeBase Shape Uniqueness)]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
num_param_idents [Params (TypeBase Shape Uniqueness)]
valueparams'
            types_grouped_by_params :: [[[Tree (TypeBase Shape Uniqueness)]]]
types_grouped_by_params = [Int]
-> [[Tree (TypeBase Shape Uniqueness)]]
-> [[[Tree (TypeBase Shape Uniqueness)]]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
num_param_idents [[Tree (TypeBase Shape Uniqueness)]]
params_ts

        [FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a
m (Params (TypeBase Shape Uniqueness)
forall {shape} {u}. [Param (TypeBase shape u)]
certparams Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
forall a. [a] -> [a] -> [a]
++ Params (TypeBase Shape Uniqueness)
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params) ([[Tree (FParam SOACS)]] -> InternaliseM a)
-> [[Tree (FParam SOACS)]] -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
          ([[Tree (TypeBase Shape Uniqueness)]]
 -> [Params (TypeBase Shape Uniqueness)]
 -> [Tree (Param (TypeBase Shape Uniqueness))])
-> [[[Tree (TypeBase Shape Uniqueness)]]]
-> [[Params (TypeBase Shape Uniqueness)]]
-> [[Tree (Param (TypeBase Shape Uniqueness))]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [[Tree (TypeBase Shape Uniqueness)]]
-> [Params (TypeBase Shape Uniqueness)]
-> [Tree (Param (TypeBase Shape Uniqueness))]
[[Tree (TypeBase Shape Uniqueness)]]
-> [[FParam SOACS]] -> [Tree (FParam SOACS)]
chunkValues [[[Tree (TypeBase Shape Uniqueness)]]]
types_grouped_by_params [[Params (TypeBase Shape Uniqueness)]]
values_grouped_by_params
  where
    fixAccParams :: [Param (TypeBase shape u)]
-> ([Param (TypeBase shape u)], [Param (TypeBase shape u)])
fixAccParams [Param (TypeBase shape u)]
ps =
      ([Maybe (Param (TypeBase shape u))] -> [Param (TypeBase shape u)])
-> ([Maybe (Param (TypeBase shape u))], [Param (TypeBase shape u)])
-> ([Param (TypeBase shape u)], [Param (TypeBase shape u)])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [Maybe (Param (TypeBase shape u))] -> [Param (TypeBase shape u)]
forall a. [Maybe a] -> [a]
catMaybes (([Maybe (Param (TypeBase shape u))], [Param (TypeBase shape u)])
 -> ([Param (TypeBase shape u)], [Param (TypeBase shape u)]))
-> ([Maybe (Param (TypeBase shape u))], [Param (TypeBase shape u)])
-> ([Param (TypeBase shape u)], [Param (TypeBase shape u)])
forall a b. (a -> b) -> a -> b
$ [(Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))]
-> ([Maybe (Param (TypeBase shape u))], [Param (TypeBase shape u)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))]
 -> ([Maybe (Param (TypeBase shape u))],
     [Param (TypeBase shape u)]))
-> [(Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))]
-> ([Maybe (Param (TypeBase shape u))], [Param (TypeBase shape u)])
forall a b. (a -> b) -> a -> b
$ (Param (TypeBase shape u)
 -> (Maybe (Param (TypeBase shape u)), Param (TypeBase shape u)))
-> [Param (TypeBase shape u)]
-> [(Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase shape u)
-> (Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))
forall {shape} {u} {shape} {u}.
Param (TypeBase shape u)
-> (Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))
fixAccParam [Param (TypeBase shape u)]
ps
    fixAccParam :: Param (TypeBase shape u)
-> (Maybe (Param (TypeBase shape u)), Param (TypeBase shape u))
fixAccParam (I.Param Attrs
attrs VName
pv (I.Acc VName
acc Shape
ispace [Type]
ts u
u)) =
      ( Param (TypeBase shape u) -> Maybe (Param (TypeBase shape u))
forall a. a -> Maybe a
Just (Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
attrs VName
acc (TypeBase shape u -> Param (TypeBase shape u))
-> TypeBase shape u -> Param (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit),
        Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
attrs VName
pv (VName -> Shape -> [Type] -> u -> TypeBase shape u
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
I.Acc VName
acc Shape
ispace [Type]
ts u
u)
      )
    fixAccParam Param (TypeBase shape u)
p = (Maybe (Param (TypeBase shape u))
forall a. Maybe a
Nothing, Param (TypeBase shape u)
p)

    chunkValues ::
      [[Tree (I.TypeBase I.Shape Uniqueness)]] ->
      [[I.FParam I.SOACS]] ->
      [Tree (I.FParam I.SOACS)]
    chunkValues :: [[Tree (TypeBase Shape Uniqueness)]]
-> [[FParam SOACS]] -> [Tree (FParam SOACS)]
chunkValues [[Tree (TypeBase Shape Uniqueness)]]
tss [[FParam SOACS]]
vss =
      [[Tree (FParam SOACS)]] -> [Tree (FParam SOACS)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Tree (FParam SOACS)]] -> [Tree (FParam SOACS)])
-> [[Tree (FParam SOACS)]] -> [Tree (FParam SOACS)]
forall a b. (a -> b) -> a -> b
$ ([Tree (TypeBase Shape Uniqueness)]
 -> Params (TypeBase Shape Uniqueness)
 -> [Tree (Param (TypeBase Shape Uniqueness))])
-> [[Tree (TypeBase Shape Uniqueness)]]
-> [Params (TypeBase Shape Uniqueness)]
-> [[Tree (Param (TypeBase Shape Uniqueness))]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [Tree (TypeBase Shape Uniqueness)]
-> Params (TypeBase Shape Uniqueness)
-> [Tree (Param (TypeBase Shape Uniqueness))]
forall {a} {b}. [Tree a] -> [b] -> [Tree b]
f [[Tree (TypeBase Shape Uniqueness)]]
tss [Params (TypeBase Shape Uniqueness)]
[[FParam SOACS]]
vss
      where
        f :: [Tree a] -> [b] -> [Tree b]
f [Tree a]
ts [b]
vs = (Tree a -> [b] -> Tree b) -> [Tree a] -> [[b]] -> [Tree b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Tree a -> [b] -> Tree b
forall a b. Tree a -> [b] -> Tree b
treeLike [Tree a]
ts ([Int] -> [b] -> [[b]]
forall a. [Int] -> [a] -> [[a]]
chunks ((Tree a -> Int) -> [Tree a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Tree a -> Int
forall a. Free [] a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
ts) [b]
vs)

bindingLoopParams ::
  [E.TypeParam] ->
  E.Pat E.ParamType ->
  [I.Type] ->
  ([I.FParam I.SOACS] -> [I.FParam I.SOACS] -> InternaliseM a) ->
  InternaliseM a
bindingLoopParams :: forall a.
[TypeParam]
-> Pat ParamType
-> [Type]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParam]
tparams Pat ParamType
pat [Type]
ts [FParam SOACS] -> [FParam SOACS] -> InternaliseM a
m = do
  [(Ident ParamType, [AttrInfo VName])]
pat_idents <- Pat ParamType -> InternaliseM [(Ident ParamType, [AttrInfo VName])]
forall (m :: * -> *) u.
MonadFreshNames m =>
Pat (TypeBase Size u)
-> m [(Ident (TypeBase Size u), [AttrInfo VName])]
flattenPat Pat ParamType
pat
  [TypeBase Shape Uniqueness]
pat_ts <- ParamType -> [Type] -> InternaliseM [TypeBase Shape Uniqueness]
forall shape u.
ParamType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape Uniqueness]
internaliseLoopParamType (Pat ParamType -> ParamType
forall d u. Pat (TypeBase d u) -> TypeBase d u
E.patternType Pat ParamType
pat) [Type]
ts

  let shape_params :: [Param (TypeBase shape u)]
shape_params = [Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
forall a. Monoid a => a
mempty VName
v (TypeBase shape u -> Param (TypeBase shape u))
-> TypeBase shape u -> Param (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.int64 | E.TypeParamDim VName
v SrcLoc
_ <- [TypeParam]
tparams]
      shape_subst :: Map VName [SubExp]
shape_subst = [(VName, [SubExp])] -> Map VName [SubExp]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Param (TypeBase Any Any) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Any Any)
p, [VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Any Any) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Any Any)
p]) | Param (TypeBase Any Any)
p <- [Param (TypeBase Any Any)]
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params]

  [(Ident ParamType, [AttrInfo VName])]
-> [TypeBase Shape Uniqueness]
-> ([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
-> InternaliseM a
forall t a.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> ([Params t] -> InternaliseM a) -> InternaliseM a
bindingFlatPat [(Ident ParamType, [AttrInfo VName])]
pat_idents [TypeBase Shape Uniqueness]
pat_ts (([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
 -> InternaliseM a)
-> ([Params (TypeBase Shape Uniqueness)] -> InternaliseM a)
-> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \[Params (TypeBase Shape Uniqueness)]
valueparams ->
    Scope SOACS -> InternaliseM a -> InternaliseM a
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
I.localScope (Params (TypeBase Shape Uniqueness) -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
I.scopeOfFParams (Params (TypeBase Shape Uniqueness) -> Scope SOACS)
-> Params (TypeBase Shape Uniqueness) -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ Params (TypeBase Shape Uniqueness)
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
-> Params (TypeBase Shape Uniqueness)
forall a. [a] -> [a] -> [a]
++ [Params (TypeBase Shape Uniqueness)]
-> Params (TypeBase Shape Uniqueness)
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Params (TypeBase Shape Uniqueness)]
valueparams) (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
      Map VName [SubExp] -> InternaliseM a -> InternaliseM a
forall a. Map VName [SubExp] -> InternaliseM a -> InternaliseM a
substitutingVars Map VName [SubExp]
shape_subst (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
        [FParam SOACS] -> [FParam SOACS] -> InternaliseM a
m Params (TypeBase Shape Uniqueness)
[FParam SOACS]
forall {shape} {u}. [Param (TypeBase shape u)]
shape_params ([Params (TypeBase Shape Uniqueness)]
-> Params (TypeBase Shape Uniqueness)
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Params (TypeBase Shape Uniqueness)]
valueparams)

bindingLambdaParams ::
  [E.Pat E.ParamType] ->
  [I.Type] ->
  ([I.LParam I.SOACS] -> InternaliseM a) ->
  InternaliseM a
bindingLambdaParams :: forall a.
[Pat ParamType]
-> [Type] -> ([LParam SOACS] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [Pat ParamType]
params [Type]
ts [LParam SOACS] -> InternaliseM a
m = do
  [(Ident ParamType, [AttrInfo VName])]
params_idents <- [[(Ident ParamType, [AttrInfo VName])]]
-> [(Ident ParamType, [AttrInfo VName])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Ident ParamType, [AttrInfo VName])]]
 -> [(Ident ParamType, [AttrInfo VName])])
-> InternaliseM [[(Ident ParamType, [AttrInfo VName])]]
-> InternaliseM [(Ident ParamType, [AttrInfo VName])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat ParamType
 -> InternaliseM [(Ident ParamType, [AttrInfo VName])])
-> [Pat ParamType]
-> InternaliseM [[(Ident ParamType, [AttrInfo VName])]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> InternaliseM [(Ident ParamType, [AttrInfo VName])]
forall (m :: * -> *) u.
MonadFreshNames m =>
Pat (TypeBase Size u)
-> m [(Ident (TypeBase Size u), [AttrInfo VName])]
flattenPat [Pat ParamType]
params

  [(Ident ParamType, [AttrInfo VName])]
-> [Type] -> ([Params Type] -> InternaliseM a) -> InternaliseM a
forall t a.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> ([Params t] -> InternaliseM a) -> InternaliseM a
bindingFlatPat [(Ident ParamType, [AttrInfo VName])]
params_idents [Type]
ts (([Params Type] -> InternaliseM a) -> InternaliseM a)
-> ([Params Type] -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \[Params Type]
params' ->
    Scope SOACS -> InternaliseM a -> InternaliseM a
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
I.localScope (Params Type -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
I.scopeOfLParams (Params Type -> Scope SOACS) -> Params Type -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ [Params Type] -> Params Type
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Params Type]
params') (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
      [LParam SOACS] -> InternaliseM a
m ([Params Type] -> Params Type
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Params Type]
params')

type Params t = [I.Param t]

processFlatPat ::
  (Show t) =>
  [(E.Ident ParamType, [E.AttrInfo VName])] ->
  [t] ->
  InternaliseM ([Params t], VarSubsts)
processFlatPat :: forall t.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> InternaliseM ([Params t], Map VName [SubExp])
processFlatPat [(Ident ParamType, [AttrInfo VName])]
x [t]
y = [([Param t], (VName, [SubExp]))]
-> [(Ident ParamType, [AttrInfo VName])]
-> [t]
-> InternaliseM ([[Param t]], Map VName [SubExp])
forall {dec}.
[([Param dec], (VName, [SubExp]))]
-> [(Ident ParamType, [AttrInfo VName])]
-> [dec]
-> InternaliseM ([[Param dec]], Map VName [SubExp])
processFlatPat' [] [(Ident ParamType, [AttrInfo VName])]
x [t]
y
  where
    processFlatPat' :: [([Param dec], (VName, [SubExp]))]
-> [(Ident ParamType, [AttrInfo VName])]
-> [dec]
-> InternaliseM ([[Param dec]], Map VName [SubExp])
processFlatPat' [([Param dec], (VName, [SubExp]))]
pat [] [dec]
_ = do
      let ([[Param dec]]
vs, [(VName, [SubExp])]
substs) = [([Param dec], (VName, [SubExp]))]
-> ([[Param dec]], [(VName, [SubExp])])
forall a b. [(a, b)] -> ([a], [b])
unzip [([Param dec], (VName, [SubExp]))]
pat
      ([[Param dec]], Map VName [SubExp])
-> InternaliseM ([[Param dec]], Map VName [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[Param dec]] -> [[Param dec]]
forall a. [a] -> [a]
reverse [[Param dec]]
vs, [(VName, [SubExp])] -> Map VName [SubExp]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, [SubExp])]
substs)
    processFlatPat' [([Param dec], (VName, [SubExp]))]
pat ((Ident ParamType
p, [AttrInfo VName]
attrs) : [(Ident ParamType, [AttrInfo VName])]
rest) [dec]
ts = do
      Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
      ([Param dec]
ps, [dec]
rest_ts) <- Attrs -> [dec] -> [VName] -> ([Param dec], [dec])
forall {dec}. Attrs -> [dec] -> [VName] -> ([Param dec], [dec])
handleMapping Attrs
attrs' [dec]
ts ([VName] -> ([Param dec], [dec]))
-> InternaliseM [VName] -> InternaliseM ([Param dec], [dec])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ident ParamType -> InternaliseM [VName]
internaliseBindee Ident ParamType
p
      [([Param dec], (VName, [SubExp]))]
-> [(Ident ParamType, [AttrInfo VName])]
-> [dec]
-> InternaliseM ([[Param dec]], Map VName [SubExp])
processFlatPat'
        (([Param dec]
ps, (Ident ParamType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName Ident ParamType
p, (Param dec -> SubExp) -> [Param dec] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param dec -> VName) -> Param dec -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param dec -> VName
forall dec. Param dec -> VName
I.paramName) [Param dec]
ps)) ([Param dec], (VName, [SubExp]))
-> [([Param dec], (VName, [SubExp]))]
-> [([Param dec], (VName, [SubExp]))]
forall a. a -> [a] -> [a]
: [([Param dec], (VName, [SubExp]))]
pat)
        [(Ident ParamType, [AttrInfo VName])]
rest
        [dec]
rest_ts

    handleMapping :: Attrs -> [dec] -> [VName] -> ([Param dec], [dec])
handleMapping Attrs
_ [dec]
ts [] =
      ([], [dec]
ts)
    handleMapping Attrs
attrs (dec
t : [dec]
ts) (VName
r : [VName]
rs) =
      let ([Param dec]
ps, [dec]
ts') = Attrs -> [dec] -> [VName] -> ([Param dec], [dec])
handleMapping Attrs
attrs [dec]
ts [VName]
rs
       in (Attrs -> VName -> dec -> Param dec
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
attrs VName
r dec
t Param dec -> [Param dec] -> [Param dec]
forall a. a -> [a] -> [a]
: [Param dec]
ps, [dec]
ts')
    handleMapping Attrs
_ [] [VName]
_ =
      [Char] -> ([Param dec], [dec])
forall a. HasCallStack => [Char] -> a
error ([Char] -> ([Param dec], [dec])) -> [Char] -> ([Param dec], [dec])
forall a b. (a -> b) -> a -> b
$ [Char]
"handleMapping: insufficient identifiers in pattern.\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ([(Ident ParamType, [AttrInfo VName])], [t]) -> [Char]
forall a. Show a => a -> [Char]
show ([(Ident ParamType, [AttrInfo VName])]
x, [t]
y)

    internaliseBindee :: E.Ident E.ParamType -> InternaliseM [VName]
    internaliseBindee :: Ident ParamType -> InternaliseM [VName]
internaliseBindee Ident ParamType
bindee = do
      let name :: VName
name = Ident ParamType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName Ident ParamType
bindee
      case ParamType -> Int
forall als. TypeBase Size als -> Int
internalisedTypeSize (ParamType -> Int) -> ParamType -> Int
forall a b. (a -> b) -> a -> b
$ Info ParamType -> ParamType
forall a. Info a -> a
E.unInfo (Info ParamType -> ParamType) -> Info ParamType -> ParamType
forall a b. (a -> b) -> a -> b
$ Ident ParamType -> Info ParamType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
E.identType Ident ParamType
bindee of
        Int
1 -> [VName] -> InternaliseM [VName]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName
name]
        Int
n -> Int -> InternaliseM VName -> InternaliseM [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (InternaliseM VName -> InternaliseM [VName])
-> InternaliseM VName -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> InternaliseM VName) -> [Char] -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString VName
name

bindingFlatPat ::
  (Show t) =>
  [(E.Ident E.ParamType, [E.AttrInfo VName])] ->
  [t] ->
  ([Params t] -> InternaliseM a) ->
  InternaliseM a
bindingFlatPat :: forall t a.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> ([Params t] -> InternaliseM a) -> InternaliseM a
bindingFlatPat [(Ident ParamType, [AttrInfo VName])]
idents [t]
ts [Params t] -> InternaliseM a
m = do
  ([Params t]
ps, Map VName [SubExp]
substs) <- [(Ident ParamType, [AttrInfo VName])]
-> [t] -> InternaliseM ([Params t], Map VName [SubExp])
forall t.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> InternaliseM ([Params t], Map VName [SubExp])
processFlatPat [(Ident ParamType, [AttrInfo VName])]
idents [t]
ts
  (InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall a.
(InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\InternaliseEnv
env -> InternaliseEnv
env {envSubsts :: Map VName [SubExp]
envSubsts = Map VName [SubExp]
substs Map VName [SubExp] -> Map VName [SubExp] -> Map VName [SubExp]
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` InternaliseEnv -> Map VName [SubExp]
envSubsts InternaliseEnv
env}) (InternaliseM a -> InternaliseM a)
-> InternaliseM a -> InternaliseM a
forall a b. (a -> b) -> a -> b
$
    [Params t] -> InternaliseM a
m [Params t]
ps

-- | Flatten a pattern.  Returns a list of identifiers.
flattenPat :: (MonadFreshNames m) => E.Pat (TypeBase Size u) -> m [(E.Ident (TypeBase Size u), [E.AttrInfo VName])]
flattenPat :: forall (m :: * -> *) u.
MonadFreshNames m =>
Pat (TypeBase Size u)
-> m [(Ident (TypeBase Size u), [AttrInfo VName])]
flattenPat = PatBase Info VName (TypeBase Size u)
-> m [(IdentBase Info VName (TypeBase Size u), [AttrInfo VName])]
forall {f :: * -> *} {dim} {u}.
MonadFreshNames f =>
PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat'
  where
    flattenPat' :: PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (E.PatParens PatBase Info VName (TypeBase dim u)
p SrcLoc
_) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' PatBase Info VName (TypeBase dim u)
p
    flattenPat' (E.PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase dim u)
p SrcLoc
_) =
      ((IdentBase Info VName (TypeBase dim u), [AttrInfo VName])
 -> (IdentBase Info VName (TypeBase dim u), [AttrInfo VName]))
-> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
-> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall a b. (a -> b) -> [a] -> [b]
map (([AttrInfo VName] -> [AttrInfo VName])
-> (IdentBase Info VName (TypeBase dim u), [AttrInfo VName])
-> (IdentBase Info VName (TypeBase dim u), [AttrInfo VName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (AttrInfo VName
attr :)) ([(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
 -> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' PatBase Info VName (TypeBase dim u)
p
    flattenPat' (E.Wildcard Info (TypeBase dim u)
t SrcLoc
loc) = do
      VName
name <- [Char] -> f VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"nameless"
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (PatBase Info VName (TypeBase dim u)
 -> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall a b. (a -> b) -> a -> b
$ VName
-> Info (TypeBase dim u)
-> SrcLoc
-> PatBase Info VName (TypeBase dim u)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
E.Id VName
name Info (TypeBase dim u)
t SrcLoc
loc
    flattenPat' (E.Id VName
v (Info TypeBase dim u
t) SrcLoc
loc) =
      [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(VName
-> Info (TypeBase dim u)
-> SrcLoc
-> IdentBase Info VName (TypeBase dim u)
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
E.Ident VName
v (TypeBase dim u -> Info (TypeBase dim u)
forall a. a -> Info a
Info TypeBase dim u
t) SrcLoc
loc, [AttrInfo VName]
forall a. Monoid a => a
mempty)]
    flattenPat' (E.TuplePat [] SrcLoc
loc) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (Info (TypeBase dim u)
-> SrcLoc -> PatBase Info VName (TypeBase dim u)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (TypeBase dim u -> Info (TypeBase dim u)
forall a. a -> Info a
Info (TypeBase dim u -> Info (TypeBase dim u))
-> TypeBase dim u -> Info (TypeBase dim u)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim u) -> ScalarTypeBase dim u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record Map Name (TypeBase dim u)
forall a. Monoid a => a
mempty) SrcLoc
loc)
    flattenPat' (E.RecordPat [] SrcLoc
loc) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (Info (TypeBase dim u)
-> SrcLoc -> PatBase Info VName (TypeBase dim u)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (TypeBase dim u -> Info (TypeBase dim u)
forall a. a -> Info a
Info (TypeBase dim u -> Info (TypeBase dim u))
-> TypeBase dim u -> Info (TypeBase dim u)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim u) -> ScalarTypeBase dim u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record Map Name (TypeBase dim u)
forall a. Monoid a => a
mempty) SrcLoc
loc)
    flattenPat' (E.TuplePat [PatBase Info VName (TypeBase dim u)]
pats SrcLoc
_) =
      [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
-> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
 -> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> f [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase Info VName (TypeBase dim u)
 -> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> [PatBase Info VName (TypeBase dim u)]
-> f [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' [PatBase Info VName (TypeBase dim u)]
pats
    flattenPat' (E.RecordPat [(Name, PatBase Info VName (TypeBase dim u))]
fs SrcLoc
loc) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (PatBase Info VName (TypeBase dim u)
 -> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName (TypeBase dim u)]
-> SrcLoc -> PatBase Info VName (TypeBase dim u)
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
E.TuplePat (((Name, PatBase Info VName (TypeBase dim u))
 -> PatBase Info VName (TypeBase dim u))
-> [(Name, PatBase Info VName (TypeBase dim u))]
-> [PatBase Info VName (TypeBase dim u)]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatBase Info VName (TypeBase dim u))
-> PatBase Info VName (TypeBase dim u)
forall a b. (a, b) -> b
snd ([(Name, PatBase Info VName (TypeBase dim u))]
 -> [PatBase Info VName (TypeBase dim u)])
-> [(Name, PatBase Info VName (TypeBase dim u))]
-> [PatBase Info VName (TypeBase dim u)]
forall a b. (a -> b) -> a -> b
$ Map Name (PatBase Info VName (TypeBase dim u))
-> [(Name, PatBase Info VName (TypeBase dim u))]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name (PatBase Info VName (TypeBase dim u))
 -> [(Name, PatBase Info VName (TypeBase dim u))])
-> Map Name (PatBase Info VName (TypeBase dim u))
-> [(Name, PatBase Info VName (TypeBase dim u))]
forall a b. (a -> b) -> a -> b
$ [(Name, PatBase Info VName (TypeBase dim u))]
-> Map Name (PatBase Info VName (TypeBase dim u))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info VName (TypeBase dim u))]
fs) SrcLoc
loc
    flattenPat' (E.PatAscription PatBase Info VName (TypeBase dim u)
p TypeExp Info VName
_ SrcLoc
_) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' PatBase Info VName (TypeBase dim u)
p
    flattenPat' (E.PatLit PatLit
_ Info (TypeBase dim u)
t SrcLoc
loc) =
      PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' (PatBase Info VName (TypeBase dim u)
 -> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall a b. (a -> b) -> a -> b
$ Info (TypeBase dim u)
-> SrcLoc -> PatBase Info VName (TypeBase dim u)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard Info (TypeBase dim u)
t SrcLoc
loc
    flattenPat' (E.PatConstr Name
_ Info (TypeBase dim u)
_ [PatBase Info VName (TypeBase dim u)]
ps SrcLoc
_) =
      [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
-> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
 -> [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> f [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase Info VName (TypeBase dim u)
 -> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])])
-> [PatBase Info VName (TypeBase dim u)]
-> f [[(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PatBase Info VName (TypeBase dim u)
-> f [(IdentBase Info VName (TypeBase dim u), [AttrInfo VName])]
flattenPat' [PatBase Info VName (TypeBase dim u)]
ps

stmPat ::
  E.Pat E.ParamType ->
  [I.Type] ->
  ([VName] -> InternaliseM a) ->
  InternaliseM a
stmPat :: forall a.
Pat ParamType
-> [Type] -> ([VName] -> InternaliseM a) -> InternaliseM a
stmPat Pat ParamType
pat [Type]
ts [VName] -> InternaliseM a
m = do
  [(Ident ParamType, [AttrInfo VName])]
pat' <- Pat ParamType -> InternaliseM [(Ident ParamType, [AttrInfo VName])]
forall (m :: * -> *) u.
MonadFreshNames m =>
Pat (TypeBase Size u)
-> m [(Ident (TypeBase Size u), [AttrInfo VName])]
flattenPat Pat ParamType
pat
  [(Ident ParamType, [AttrInfo VName])]
-> [Type] -> ([Params Type] -> InternaliseM a) -> InternaliseM a
forall t a.
Show t =>
[(Ident ParamType, [AttrInfo VName])]
-> [t] -> ([Params t] -> InternaliseM a) -> InternaliseM a
bindingFlatPat [(Ident ParamType, [AttrInfo VName])]
pat' [Type]
ts (([Params Type] -> InternaliseM a) -> InternaliseM a)
-> ([Params Type] -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ [VName] -> InternaliseM a
m ([VName] -> InternaliseM a)
-> ([Params Type] -> [VName]) -> [Params Type] -> InternaliseM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type -> VName) -> Params Type -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
I.paramName (Params Type -> [VName])
-> ([Params Type] -> Params Type) -> [Params Type] -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Params Type] -> Params Type
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat