{-# LANGUAGE TypeFamilies #-}

module Futhark.Analysis.HORep.MapNest
  ( Nesting (..),
    MapNest (..),
    typeOf,
    params,
    inputs,
    setInputs,
    fromSOAC,
    toSOAC,
  )
where

import Data.List (find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.HORep.SOAC (SOAC)
import Futhark.Analysis.HORep.SOAC qualified as SOAC
import Futhark.Construct
import Futhark.IR hiding (typeOf)
import Futhark.IR.SOACS.SOAC qualified as Futhark
import Futhark.Transform.Substitute

data Nesting rep = Nesting
  { forall {k} (rep :: k). Nesting rep -> [VName]
nestingParamNames :: [VName],
    forall {k} (rep :: k). Nesting rep -> [VName]
nestingResult :: [VName],
    forall {k} (rep :: k). Nesting rep -> [Type]
nestingReturnType :: [Type],
    forall {k} (rep :: k). Nesting rep -> SubExp
nestingWidth :: SubExp
  }
  deriving (Nesting rep -> Nesting rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
/= :: Nesting rep -> Nesting rep -> Bool
$c/= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
== :: Nesting rep -> Nesting rep -> Bool
$c== :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
Eq, Nesting rep -> Nesting rep -> Bool
Nesting rep -> Nesting rep -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k (rep :: k). Eq (Nesting rep)
forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
forall k (rep :: k). Nesting rep -> Nesting rep -> Ordering
forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
min :: Nesting rep -> Nesting rep -> Nesting rep
$cmin :: forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
max :: Nesting rep -> Nesting rep -> Nesting rep
$cmax :: forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
>= :: Nesting rep -> Nesting rep -> Bool
$c>= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
> :: Nesting rep -> Nesting rep -> Bool
$c> :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
<= :: Nesting rep -> Nesting rep -> Bool
$c<= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
< :: Nesting rep -> Nesting rep -> Bool
$c< :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
compare :: Nesting rep -> Nesting rep -> Ordering
$ccompare :: forall k (rep :: k). Nesting rep -> Nesting rep -> Ordering
Ord, Int -> Nesting rep -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). Int -> Nesting rep -> ShowS
forall k (rep :: k). [Nesting rep] -> ShowS
forall k (rep :: k). Nesting rep -> String
showList :: [Nesting rep] -> ShowS
$cshowList :: forall k (rep :: k). [Nesting rep] -> ShowS
show :: Nesting rep -> String
$cshow :: forall k (rep :: k). Nesting rep -> String
showsPrec :: Int -> Nesting rep -> ShowS
$cshowsPrec :: forall k (rep :: k). Int -> Nesting rep -> ShowS
Show)

data MapNest rep = MapNest SubExp (Lambda rep) [Nesting rep] [SOAC.Input]
  deriving (Int -> MapNest rep -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). RepTypes rep => Int -> MapNest rep -> ShowS
forall k (rep :: k). RepTypes rep => [MapNest rep] -> ShowS
forall k (rep :: k). RepTypes rep => MapNest rep -> String
showList :: [MapNest rep] -> ShowS
$cshowList :: forall k (rep :: k). RepTypes rep => [MapNest rep] -> ShowS
show :: MapNest rep -> String
$cshow :: forall k (rep :: k). RepTypes rep => MapNest rep -> String
showsPrec :: Int -> MapNest rep -> ShowS
$cshowsPrec :: forall k (rep :: k). RepTypes rep => Int -> MapNest rep -> ShowS
Show)

typeOf :: MapNest rep -> [Type]
typeOf :: forall {k} (rep :: k). MapNest rep -> [Type]
typeOf (MapNest SubExp
w Lambda rep
lam [] [Input]
_) =
  forall a b. (a -> b) -> [a] -> [b]
map (forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w) forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam
typeOf (MapNest SubExp
w Lambda rep
_ (Nesting rep
nest : [Nesting rep]
_) [Input]
_) =
  forall a b. (a -> b) -> [a] -> [b]
map (forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w) forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Nesting rep -> [Type]
nestingReturnType Nesting rep
nest

params :: MapNest rep -> [VName]
params :: forall {k} (rep :: k). MapNest rep -> [VName]
params (MapNest SubExp
_ Lambda rep
lam [] [Input]
_) =
  forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
params (MapNest SubExp
_ Lambda rep
_ (Nesting rep
nest : [Nesting rep]
_) [Input]
_) =
  forall {k} (rep :: k). Nesting rep -> [VName]
nestingParamNames Nesting rep
nest

inputs :: MapNest rep -> [SOAC.Input]
inputs :: forall {k} (rep :: k). MapNest rep -> [Input]
inputs (MapNest SubExp
_ Lambda rep
_ [Nesting rep]
_ [Input]
inps) = [Input]
inps

setInputs :: [SOAC.Input] -> MapNest rep -> MapNest rep
setInputs :: forall {k} (rep :: k). [Input] -> MapNest rep -> MapNest rep
setInputs [] (MapNest SubExp
w Lambda rep
body [Nesting rep]
ns [Input]
_) = forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body [Nesting rep]
ns []
setInputs (Input
inp : [Input]
inps) (MapNest SubExp
_ Lambda rep
body [Nesting rep]
ns [Input]
_) = forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body [Nesting rep]
ns' (Input
inp forall a. a -> [a] -> [a]
: [Input]
inps)
  where
    w :: SubExp
w = forall u. Int -> TypeBase (ShapeBase SubExp) u -> SubExp
arraySize Int
0 forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
    ws :: [SubExp]
ws = forall a. Int -> [a] -> [a]
drop Int
1 forall a b. (a -> b) -> a -> b
$ forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
    ns' :: [Nesting rep]
ns' = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {k} {k} {rep :: k} {rep :: k}.
Nesting rep -> SubExp -> Nesting rep
setDepth [Nesting rep]
ns [SubExp]
ws
    setDepth :: Nesting rep -> SubExp -> Nesting rep
setDepth Nesting rep
n SubExp
nw = Nesting rep
n {nestingWidth :: SubExp
nestingWidth = SubExp
nw}

fromSOAC ::
  ( Buildable rep,
    MonadFreshNames m,
    LocalScope rep m,
    Op rep ~ Futhark.SOAC rep
  ) =>
  SOAC rep ->
  m (Maybe (MapNest rep))
fromSOAC :: forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
 Op rep ~ SOAC rep) =>
SOAC rep -> m (Maybe (MapNest rep))
fromSOAC = forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
 Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' forall a. Monoid a => a
mempty

fromSOAC' ::
  ( Buildable rep,
    MonadFreshNames m,
    LocalScope rep m,
    Op rep ~ Futhark.SOAC rep
  ) =>
  [Ident] ->
  SOAC rep ->
  m (Maybe (MapNest rep))
fromSOAC' :: forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
 Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' [Ident]
bound (SOAC.Screma SubExp
w (SOAC.ScremaForm [] [] Lambda rep
lam) [Input]
inps) = do
  Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
maybenest <- case ( forall {k} (rep :: k). Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam,
                      forall {k} (rep :: k). Body rep -> Result
bodyResult forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam
                    ) of
    ([Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
e], Result
res)
      | forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
res forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var (forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat) ->
          forall {k} (rep :: k) (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam) forall a b. (a -> b) -> a -> b
$
            forall {k} (rep :: k) (m :: * -> *).
(Op rep ~ SOAC rep, HasScope rep m) =>
Exp rep -> m (Either NotSOAC (SOAC rep))
SOAC.fromExp Exp rep
e
              forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. b -> Either a b
Right forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Pat (LetDec rep)
pat,)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
 Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' [Ident]
bound')
    ([Stm rep], Result)
_ ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right forall a. Maybe a
Nothing

  case Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
maybenest of
    -- Do we have a nested MapNest?
    Right (Just (Pat (LetDec rep)
pat, mn :: MapNest rep
mn@(MapNest SubExp
inner_w Lambda rep
body' [Nesting rep]
ns' [Input]
inps'))) -> do
      ([VName]
ps, [Input]
inps'') <-
        forall a b. [(a, b)] -> ([a], [b])
unzip
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
fixInputs
            SubExp
w
            (forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam) [Input]
inps)
            (forall a b. [a] -> [b] -> [(a, b)]
zip (forall {k} (rep :: k). MapNest rep -> [VName]
params MapNest rep
mn) [Input]
inps')
      let n' :: Nesting rep
n' =
            Nesting
              { nestingParamNames :: [VName]
nestingParamNames = [VName]
ps,
                nestingResult :: [VName]
nestingResult = forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat,
                nestingReturnType :: [Type]
nestingReturnType = forall {k} (rep :: k). MapNest rep -> [Type]
typeOf MapNest rep
mn,
                nestingWidth :: SubExp
nestingWidth = SubExp
inner_w
              }
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body' (Nesting rep
n' forall a. a -> [a] -> [a]
: [Nesting rep]
ns') [Input]
inps''
    -- No nested MapNest it seems.
    Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
_ -> do
      let isBound :: VName -> Maybe Ident
isBound VName
name
            | Just Ident
param <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName
name ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
bound =
                forall a. a -> Maybe a
Just Ident
param
            | Bool
otherwise =
                forall a. Maybe a
Nothing
          boundUsedInBody :: [Ident]
boundUsedInBody =
            forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe Ident
isBound forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam
      [Ident]
newParams <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadFreshNames m =>
ShowS -> Ident -> m Ident
newIdent' (forall a. [a] -> [a] -> [a]
++ String
"_wasfree")) [Ident]
boundUsedInBody
      let subst :: Map VName VName
subst =
            forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
              forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
boundUsedInBody) (forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
newParams)
          inps' :: [Input]
inps' =
            [Input]
inps
              forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map
                (ArrayTransform -> Input -> Input
SOAC.addTransform (Certs -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ forall d. [d] -> ShapeBase d
Shape [SubExp
w]) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> Input
SOAC.identInput)
                [Ident]
boundUsedInBody
          lam' :: Lambda rep
lam' =
            Lambda rep
lam
              { lambdaBody :: Body rep
lambdaBody =
                  forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam,
                lambdaParams :: [LParam rep]
lambdaParams =
                  forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
                    forall a. [a] -> [a] -> [a]
++ [forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty VName
name Type
t | Ident VName
name Type
t <- [Ident]
newParams]
              }
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
lam' [] [Input]
inps'
  where
    bound' :: [Ident]
bound' = [Ident]
bound forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> Ident
paramIdent (forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam)
fromSOAC' [Ident]
_ SOAC rep
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

toSOAC ::
  ( MonadFreshNames m,
    HasScope rep m,
    Buildable rep,
    BuilderOps rep,
    Op rep ~ Futhark.SOAC rep
  ) =>
  MapNest rep ->
  m (SOAC rep)
toSOAC :: forall (m :: * -> *) rep.
(MonadFreshNames m, HasScope rep m, Buildable rep, BuilderOps rep,
 Op rep ~ SOAC rep) =>
MapNest rep -> m (SOAC rep)
toSOAC (MapNest SubExp
w Lambda rep
lam [] [Input]
inps) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
SOAC.Screma SubExp
w (forall {k} (rep :: k). Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
lam) [Input]
inps
toSOAC (MapNest SubExp
w Lambda rep
lam (Nesting [VName]
npnames [VName]
nres [Type]
nrettype SubExp
nw : [Nesting rep]
ns) [Input]
inps) = do
  let nparams :: [Param Type]
nparams = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty) [VName]
npnames forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Input -> Type
SOAC.inputRowType [Input]
inps
  Body rep
body <- forall {k1} {k2} (rep :: k1) (m :: * -> *) (somerep :: k2).
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
    forall {k} (rep :: k) (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams [Param Type]
nparams) forall a b. (a -> b) -> a -> b
$ do
      forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
nres
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m)) =>
SOAC (Rep m) -> m (Exp (Rep m))
SOAC.toExp
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) rep.
(MonadFreshNames m, HasScope rep m, Buildable rep, BuilderOps rep,
 Op rep ~ SOAC rep) =>
MapNest rep -> m (SOAC rep)
toSOAC (forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
nw Lambda rep
lam [Nesting rep]
ns forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Ident -> Input
SOAC.identInput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Typed dec => Param dec -> Ident
paramIdent) [Param Type]
nparams)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
nres
  let outerlam :: Lambda rep
outerlam =
        Lambda
          { lambdaParams :: [LParam rep]
lambdaParams = [Param Type]
nparams,
            lambdaBody :: Body rep
lambdaBody = Body rep
body,
            lambdaReturnType :: [Type]
lambdaReturnType = [Type]
nrettype
          }
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
SOAC.Screma SubExp
w (forall {k} (rep :: k). Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
outerlam) [Input]
inps

fixInputs ::
  MonadFreshNames m =>
  SubExp ->
  [(VName, SOAC.Input)] ->
  [(VName, SOAC.Input)] ->
  m [(VName, SOAC.Input)]
fixInputs :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
fixInputs SubExp
w [(VName, Input)]
ourInps = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (VName, Input) -> m (VName, Input)
inspect
  where
    isParam :: a -> (a, b) -> Bool
isParam a
x (a
y, b
_) = a
x forall a. Eq a => a -> a -> Bool
== a
y

    inspect :: (VName, Input) -> m (VName, Input)
inspect (VName
_, SOAC.Input ArrayTransforms
ts VName
v Type
_)
      | Just (VName
p, Input
pInp) <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (forall {a} {b}. Eq a => a -> (a, b) -> Bool
isParam VName
v) [(VName, Input)]
ourInps = do
          let pInp' :: Input
pInp' = ArrayTransforms -> Input -> Input
SOAC.transformRows ArrayTransforms
ts Input
pInp
          VName
p' <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
p
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
p', Input
pInp')
    inspect (VName
param, SOAC.Input ArrayTransforms
ts VName
a Type
t) = do
      VName
param' <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (VName -> String
baseString VName
param forall a. [a] -> [a] -> [a]
++ String
"_rep")
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
param', ArrayTransforms -> VName -> Type -> Input
SOAC.Input (ArrayTransforms
ts ArrayTransforms -> ArrayTransform -> ArrayTransforms
SOAC.|> Certs -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate forall a. Monoid a => a
mempty (forall d. [d] -> ShapeBase d
Shape [SubExp
w])) VName
a Type
t)