{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
-- |
-- Module      : Data.Array.Accelerate.Trafo.Shrink
-- Copyright   : [2012..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- The shrinking substitution arises as a restriction of beta-reduction to cases
-- where the bound variable is used zero (dead-code elimination) or one (linear
-- inlining) times. By simplifying terms, the shrinking reduction can expose
-- opportunities for further optimisation.
--
-- TODO: replace with a linear shrinking algorithm; e.g.
--
--   * Andrew Appel & Trevor Jim, "Shrinking lambda expressions in linear time".
--
--   * Nick Benton, Andrew Kennedy, Sam Lindley and Claudio Russo, "Shrinking
--     Reductions in SML.NET"
--

module Data.Array.Accelerate.Trafo.Shrink (

  -- Shrinking
  ShrinkAcc,
  shrinkExp,
  shrinkFun,

  -- Occurrence counting
  UsesOfAcc, usesOfPreAcc, usesOfExp,

) where

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.AST.Environment
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Trafo.Substitution

import qualified Data.Array.Accelerate.Debug.Stats                  as Stats

import Control.Applicative                                          hiding ( Const )
import Data.Maybe                                                   ( isJust )
import Data.Monoid
import Data.Semigroup
import Prelude                                                      hiding ( exp, seq )


data VarsRange env =
  VarsRange !(Exists (Idx env))     -- rightmost variable
            {-# UNPACK #-} !Int     -- count
            !(Maybe RangeTuple)     -- tuple

data RangeTuple
  = RTNil
  | RTSingle
  | RTPair !RangeTuple !RangeTuple

lhsVarsRange :: LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env')
lhsVarsRange :: LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env')
lhsVarsRange LeftHandSide s v env env'
lhs = case LeftHandSide s v env env'
-> Either (env :~: env') (Exists (Idx env'))
forall (s :: * -> *) v env env'.
LeftHandSide s v env env'
-> Either (env :~: env') (Exists (Idx env'))
rightIx LeftHandSide s v env env'
lhs of
  Left env :~: env'
eq -> (env :~: env') -> Either (env :~: env') (VarsRange env')
forall a b. a -> Either a b
Left env :~: env'
eq
  Right Exists (Idx env')
ix -> let (Int
n, Maybe RangeTuple
rt) = LeftHandSide s v env env' -> (Int, Maybe RangeTuple)
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> (Int, Maybe RangeTuple)
go LeftHandSide s v env env'
lhs
              in  VarsRange env' -> Either (env :~: env') (VarsRange env')
forall a b. b -> Either a b
Right (VarsRange env' -> Either (env :~: env') (VarsRange env'))
-> VarsRange env' -> Either (env :~: env') (VarsRange env')
forall a b. (a -> b) -> a -> b
$ Exists (Idx env') -> Int -> Maybe RangeTuple -> VarsRange env'
forall env.
Exists (Idx env) -> Int -> Maybe RangeTuple -> VarsRange env
VarsRange Exists (Idx env')
ix Int
n Maybe RangeTuple
rt
  where
    rightIx :: LeftHandSide s v env env' -> Either (env :~: env') (Exists (Idx env'))
    rightIx :: LeftHandSide s v env env'
-> Either (env :~: env') (Exists (Idx env'))
rightIx (LeftHandSideWildcard TupR s v
_) = (env :~: env) -> Either (env :~: env) (Exists (Idx env'))
forall a b. a -> Either a b
Left env :~: env
forall k (a :: k). a :~: a
Refl
    rightIx (LeftHandSideSingle s v
_)   = Exists (Idx (env, v))
-> Either (env :~: env') (Exists (Idx (env, v)))
forall a b. b -> Either a b
Right (Exists (Idx (env, v))
 -> Either (env :~: env') (Exists (Idx (env, v))))
-> Exists (Idx (env, v))
-> Either (env :~: env') (Exists (Idx (env, v)))
forall a b. (a -> b) -> a -> b
$ Idx (env, v) v -> Exists (Idx (env, v))
forall (f :: * -> *) a. f a -> Exists f
Exists Idx (env, v) v
forall env t. Idx (env, t) t
ZeroIdx
    rightIx (LeftHandSidePair LeftHandSide s v1 env env'
l1 LeftHandSide s v2 env' env'
l2) = case LeftHandSide s v2 env' env'
-> Either (env' :~: env') (Exists (Idx env'))
forall (s :: * -> *) v env env'.
LeftHandSide s v env env'
-> Either (env :~: env') (Exists (Idx env'))
rightIx LeftHandSide s v2 env' env'
l2 of
      Right Exists (Idx env')
ix  -> Exists (Idx env') -> Either (env :~: env') (Exists (Idx env'))
forall a b. b -> Either a b
Right Exists (Idx env')
ix
      Left env' :~: env'
Refl -> LeftHandSide s v1 env env'
-> Either (env :~: env') (Exists (Idx env'))
forall (s :: * -> *) v env env'.
LeftHandSide s v env env'
-> Either (env :~: env') (Exists (Idx env'))
rightIx LeftHandSide s v1 env env'
l1

    go :: LeftHandSide s v env env' -> (Int, Maybe (RangeTuple))
    go :: LeftHandSide s v env env' -> (Int, Maybe RangeTuple)
go (LeftHandSideWildcard TupR s v
TupRunit)   = (Int
0,       RangeTuple -> Maybe RangeTuple
forall a. a -> Maybe a
Just RangeTuple
RTNil)
    go (LeftHandSideWildcard TupR s v
_)          = (Int
0,       Maybe RangeTuple
forall a. Maybe a
Nothing)
    go (LeftHandSideSingle s v
_)            = (Int
1,       RangeTuple -> Maybe RangeTuple
forall a. a -> Maybe a
Just RangeTuple
RTSingle)
    go (LeftHandSidePair LeftHandSide s v1 env env'
l1 LeftHandSide s v2 env' env'
l2)          = (Int
n1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n2, RangeTuple -> RangeTuple -> RangeTuple
RTPair (RangeTuple -> RangeTuple -> RangeTuple)
-> Maybe RangeTuple -> Maybe (RangeTuple -> RangeTuple)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe RangeTuple
t1 Maybe (RangeTuple -> RangeTuple)
-> Maybe RangeTuple -> Maybe RangeTuple
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe RangeTuple
t2)
      where
        (Int
n1, Maybe RangeTuple
t1) = LeftHandSide s v1 env env' -> (Int, Maybe RangeTuple)
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> (Int, Maybe RangeTuple)
go LeftHandSide s v1 env env'
l1
        (Int
n2, Maybe RangeTuple
t2) = LeftHandSide s v2 env' env' -> (Int, Maybe RangeTuple)
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> (Int, Maybe RangeTuple)
go LeftHandSide s v2 env' env'
l2

weakenVarsRange :: LeftHandSide s v env env' -> VarsRange env -> VarsRange env'
weakenVarsRange :: LeftHandSide s v env env' -> VarsRange env -> VarsRange env'
weakenVarsRange LeftHandSide s v env env'
lhs (VarsRange Exists (Idx env)
ix Int
n Maybe RangeTuple
t) = Exists (Idx env') -> Int -> Maybe RangeTuple -> VarsRange env'
forall env.
Exists (Idx env) -> Int -> Maybe RangeTuple -> VarsRange env
VarsRange (LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
go LeftHandSide s v env env'
lhs Exists (Idx env)
ix) Int
n Maybe RangeTuple
t
  where
    go :: LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
    go :: LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
go (LeftHandSideWildcard TupR s v
_) Exists (Idx env)
ix'          = Exists (Idx env)
Exists (Idx env')
ix'
    go (LeftHandSideSingle s v
_)   (Exists Idx env a
ix') = Idx (env, v) a -> Exists (Idx (env, v))
forall (f :: * -> *) a. f a -> Exists f
Exists (Idx env a -> Idx (env, v) a
forall env t s. Idx env t -> Idx (env, s) t
SuccIdx Idx env a
ix')
    go (LeftHandSidePair LeftHandSide s v1 env env'
l1 LeftHandSide s v2 env' env'
l2) Exists (Idx env)
ix'          = LeftHandSide s v2 env' env'
-> Exists (Idx env') -> Exists (Idx env')
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
go LeftHandSide s v2 env' env'
l2 (Exists (Idx env') -> Exists (Idx env'))
-> Exists (Idx env') -> Exists (Idx env')
forall a b. (a -> b) -> a -> b
$ LeftHandSide s v1 env env' -> Exists (Idx env) -> Exists (Idx env')
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env')
go LeftHandSide s v1 env env'
l1 Exists (Idx env)
ix'

matchEVarsRange :: VarsRange env -> OpenExp env aenv t -> Bool
matchEVarsRange :: VarsRange env -> OpenExp env aenv t -> Bool
matchEVarsRange (VarsRange (Exists Idx env a
first) Int
_ (Just RangeTuple
rt)) OpenExp env aenv t
expr = Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool) -> Maybe Int -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
forall env aenv t.
Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
go (Idx env a -> Int
forall env t. Idx env t -> Int
idxToInt Idx env a
first) RangeTuple
rt OpenExp env aenv t
expr
  where
    go :: Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
    go :: Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
go Int
i RangeTuple
RTNil OpenExp env aenv t
Nil = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i
    go Int
i RangeTuple
RTSingle (Evar (Var ScalarType t
_ Idx env t
ix))
      | Int -> Idx env t -> Bool
forall env t. Int -> Idx env t -> Bool
checkIdx Int
i Idx env t
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    go Int
i (RTPair RangeTuple
t1 RangeTuple
t2) (Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2)
      | Just Int
i' <- Int -> RangeTuple -> OpenExp env aenv t2 -> Maybe Int
forall env aenv t.
Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
go Int
i RangeTuple
t2 OpenExp env aenv t2
e2 = Int -> RangeTuple -> OpenExp env aenv t1 -> Maybe Int
forall env aenv t.
Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int
go Int
i' RangeTuple
t1 OpenExp env aenv t1
e1
    go Int
_ RangeTuple
_ OpenExp env aenv t
_ = Maybe Int
forall a. Maybe a
Nothing

    checkIdx :: Int -> Idx env t ->  Bool
    checkIdx :: Int -> Idx env t -> Bool
checkIdx Int
0 Idx env t
ZeroIdx = Bool
True
    checkIdx Int
i (SuccIdx Idx env t
ix) = Int -> Idx env t -> Bool
forall env t. Int -> Idx env t -> Bool
checkIdx (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Idx env t
ix
    checkIdx Int
_ Idx env t
_ = Bool
False
matchEVarsRange VarsRange env
_ OpenExp env aenv t
_ = Bool
False

varInRange :: VarsRange env -> Var s env t -> Maybe Usages
varInRange :: VarsRange env -> Var s env t -> Maybe Usages
varInRange (VarsRange (Exists Idx env a
rangeIx) Int
n Maybe RangeTuple
_) (Var s t
_ Idx env t
varIx) = case Idx env a -> Idx env t -> Maybe Int
forall env u t. Idx env u -> Idx env t -> Maybe Int
go Idx env a
rangeIx Idx env t
varIx of
    Maybe Int
Nothing -> Maybe Usages
forall a. Maybe a
Nothing
    Just Int
j  -> Usages -> Maybe Usages
forall a. a -> Maybe a
Just (Usages -> Maybe Usages) -> Usages -> Maybe Usages
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Usages
forall a. Int -> a -> [a]
replicate Int
j Bool
False Usages -> Usages -> Usages
forall a. [a] -> [a] -> [a]
++ [Bool
True] Usages -> Usages -> Usages
forall a. [a] -> [a] -> [a]
++ Int -> Bool -> Usages
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Bool
False
  where
    -- `go ix ix'` checks whether ix <= ix' with recursion, and then checks
    -- whether ix' < ix + n in go'. Returns a Just if both checks
    -- are successful, containing an integer j such that ix + j = ix'.
    go :: Idx env u -> Idx env t -> Maybe Int
    go :: Idx env u -> Idx env t -> Maybe Int
go (SuccIdx Idx env u
ix) (SuccIdx Idx env t
ix') = Idx env u -> Idx env t -> Maybe Int
forall env u t. Idx env u -> Idx env t -> Maybe Int
go Idx env u
ix Idx env t
Idx env t
ix'
    go Idx env u
ZeroIdx      Idx env t
ix'           = Idx env t -> Int -> Maybe Int
forall env t. Idx env t -> Int -> Maybe Int
go' Idx env t
ix' Int
0
    go Idx env u
_            Idx env t
ZeroIdx       = Maybe Int
forall a. Maybe a
Nothing

    go' :: Idx env t -> Int -> Maybe Int
    go' :: Idx env t -> Int -> Maybe Int
go' Idx env t
_ Int
j | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = Maybe Int
forall a. Maybe a
Nothing
    go' Idx env t
ZeroIdx       Int
j = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
j
    go' (SuccIdx Idx env t
ix') Int
j = Idx env t -> Int -> Maybe Int
forall env t. Idx env t -> Int -> Maybe Int
go' Idx env t
ix' (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- Describes how often the variables defined in a LHS are used together.
data Count
  = Impossible !Usages
      -- Cannot inline this definition. This happens when the definition
      -- declares multiple variables (the right hand side returns a tuple)
      -- and the variables are used seperately.
  | Infinity
      -- The variable is used in a loop. Inlining should only proceed if
      -- the computation is cheap.
  | Finite {-# UNPACK #-} !Int

type Usages = [Bool] -- Per variable a Boolean denoting whether that variable is used.

instance Semigroup Count where
  Impossible Usages
u1 <> :: Count -> Count -> Count
<> Impossible Usages
u2 = Usages -> Count
Impossible (Usages -> Count) -> Usages -> Count
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool -> Bool) -> Usages -> Usages -> Usages
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bool -> Bool -> Bool
(||) Usages
u1 Usages
u2
  Impossible Usages
u  <> Finite Int
0      = Usages -> Count
Impossible Usages
u
  Finite Int
0      <> Impossible Usages
u  = Usages -> Count
Impossible Usages
u
  Impossible Usages
u  <> Count
_             = Usages -> Count
Impossible (Usages -> Count) -> Usages -> Count
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> Usages -> Usages
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Bool -> Bool
forall a b. a -> b -> a
const Bool
True) Usages
u
  Count
_             <> Impossible Usages
u  = Usages -> Count
Impossible (Usages -> Count) -> Usages -> Count
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> Usages -> Usages
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Bool -> Bool
forall a b. a -> b -> a
const Bool
True) Usages
u
  Count
Infinity      <> Count
_             = Count
Infinity
  Count
_             <> Count
Infinity      = Count
Infinity
  Finite Int
a      <> Finite Int
b      = Int -> Count
Finite (Int -> Count) -> Int -> Count
forall a b. (a -> b) -> a -> b
$ Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b

instance Monoid Count where
  mempty :: Count
mempty = Int -> Count
Finite Int
0

loopCount :: Count -> Count
loopCount :: Count -> Count
loopCount (Finite Int
n) | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Count
Infinity
loopCount Count
c                  = Count
c

shrinkLhs
    :: HasCallStack
    => Count
    -> LeftHandSide s t env1 env2
    -> Maybe (Exists (LeftHandSide s t env1))
shrinkLhs :: Count
-> LeftHandSide s t env1 env2
-> Maybe (Exists (LeftHandSide s t env1))
shrinkLhs Count
_ (LeftHandSideWildcard TupR s t
_) = Maybe (Exists (LeftHandSide s t env1))
forall a. Maybe a
Nothing -- We cannot shrink this
shrinkLhs (Finite Int
0)          LeftHandSide s t env1 env2
lhs = Exists (LeftHandSide s t env1)
-> Maybe (Exists (LeftHandSide s t env1))
forall a. a -> Maybe a
Just (Exists (LeftHandSide s t env1)
 -> Maybe (Exists (LeftHandSide s t env1)))
-> Exists (LeftHandSide s t env1)
-> Maybe (Exists (LeftHandSide s t env1))
forall a b. (a -> b) -> a -> b
$ LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall (f :: * -> *) a. f a -> Exists f
Exists (LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1))
-> LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall a b. (a -> b) -> a -> b
$ TupR s t -> LeftHandSide s t env1 env1
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard (TupR s t -> LeftHandSide s t env1 env1)
-> TupR s t -> LeftHandSide s t env1 env1
forall a b. (a -> b) -> a -> b
$ LeftHandSide s t env1 env2 -> TupR s t
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> TupR s v
lhsToTupR LeftHandSide s t env1 env2
lhs -- LHS isn't used at all, replace with a wildcard
shrinkLhs (Impossible Usages
usages) LeftHandSide s t env1 env2
lhs = case Usages
-> LeftHandSide s t env1 env2
-> (Bool, Usages, Exists (LeftHandSide s t env1))
forall (s :: * -> *) t env1 env2.
HasCallStack =>
Usages
-> LeftHandSide s t env1 env2
-> (Bool, Usages, Exists (LeftHandSide s t env1))
go Usages
usages LeftHandSide s t env1 env2
lhs of
    (Bool
True , [], Exists (LeftHandSide s t env1)
lhs') -> Exists (LeftHandSide s t env1)
-> Maybe (Exists (LeftHandSide s t env1))
forall a. a -> Maybe a
Just Exists (LeftHandSide s t env1)
lhs'
    (Bool
False, [], Exists (LeftHandSide s t env1)
_   ) -> Maybe (Exists (LeftHandSide s t env1))
forall a. Maybe a
Nothing -- No variables were dropped. Thus lhs == lhs'.
    (Bool, Usages, Exists (LeftHandSide s t env1))
_                 -> String -> Maybe (Exists (LeftHandSide s t env1))
forall a. HasCallStack => String -> a
internalError String
"Mismatch in length of usages array and LHS"
  where
    go :: HasCallStack => Usages -> LeftHandSide s t env1 env2 -> (Bool, Usages, Exists (LeftHandSide s t env1))
    go :: Usages
-> LeftHandSide s t env1 env2
-> (Bool, Usages, Exists (LeftHandSide s t env1))
go Usages
us           (LeftHandSideWildcard TupR s t
tp) = (Bool
False, Usages
us, LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall (f :: * -> *) a. f a -> Exists f
Exists (LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1))
-> LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall a b. (a -> b) -> a -> b
$ TupR s t -> LeftHandSide s t env1 env1
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard TupR s t
tp)
    go (Bool
True  : Usages
us) (LeftHandSideSingle s t
tp)   = (Bool
False, Usages
us, LeftHandSide s t env1 (env1, t) -> Exists (LeftHandSide s t env1)
forall (f :: * -> *) a. f a -> Exists f
Exists (LeftHandSide s t env1 (env1, t) -> Exists (LeftHandSide s t env1))
-> LeftHandSide s t env1 (env1, t)
-> Exists (LeftHandSide s t env1)
forall a b. (a -> b) -> a -> b
$ s t -> LeftHandSide s t env1 (env1, t)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle s t
tp)
    go (Bool
False : Usages
us) (LeftHandSideSingle s t
tp)   = (Bool
True , Usages
us, LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall (f :: * -> *) a. f a -> Exists f
Exists (LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1))
-> LeftHandSide s t env1 env1 -> Exists (LeftHandSide s t env1)
forall a b. (a -> b) -> a -> b
$ TupR s t -> LeftHandSide s t env1 env1
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard (TupR s t -> LeftHandSide s t env1 env1)
-> TupR s t -> LeftHandSide s t env1 env1
forall a b. (a -> b) -> a -> b
$ s t -> TupR s t
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle s t
tp)
    go Usages
us           (LeftHandSidePair LeftHandSide s v1 env1 env'
l1 LeftHandSide s v2 env' env2
l2)
      | (Bool
c2, Usages
us' , Exists LeftHandSide s v2 env' a
l2') <- Usages
-> LeftHandSide s v2 env' env2
-> (Bool, Usages, Exists (LeftHandSide s v2 env'))
forall (s :: * -> *) t env1 env2.
HasCallStack =>
Usages
-> LeftHandSide s t env1 env2
-> (Bool, Usages, Exists (LeftHandSide s t env1))
go Usages
us  LeftHandSide s v2 env' env2
l2
      , (Bool
c1, Usages
us'', Exists LeftHandSide s v1 env1 a
l1') <- Usages
-> LeftHandSide s v1 env1 env'
-> (Bool, Usages, Exists (LeftHandSide s v1 env1))
forall (s :: * -> *) t env1 env2.
HasCallStack =>
Usages
-> LeftHandSide s t env1 env2
-> (Bool, Usages, Exists (LeftHandSide s t env1))
go Usages
us' LeftHandSide s v1 env1 env'
l1
      , Exists LeftHandSide s v2 a a
l2'' <- LeftHandSide s v2 env' a -> Exists (LeftHandSide s v2 a)
forall (s :: * -> *) t aenv1 aenv1' aenv2.
LeftHandSide s t aenv1 aenv1' -> Exists (LeftHandSide s t aenv2)
rebuildLHS LeftHandSide s v2 env' a
l2'
      = let
          lhs' :: LeftHandSide s (v1, v2) env1 a
lhs'
            | LeftHandSideWildcard TupR s v1
t1 <- LeftHandSide s v1 env1 a
l1'
            , LeftHandSideWildcard TupR s v2
t2 <- LeftHandSide s v2 a a
l2'' = TupR s (v1, v2) -> LeftHandSide s (v1, v2) a a
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard (TupR s (v1, v2) -> LeftHandSide s (v1, v2) a a)
-> TupR s (v1, v2) -> LeftHandSide s (v1, v2) a a
forall a b. (a -> b) -> a -> b
$ TupR s v1 -> TupR s v2 -> TupR s (v1, v2)
forall (s :: * -> *) a b. TupR s a -> TupR s b -> TupR s (a, b)
TupRpair TupR s v1
t1 TupR s v2
t2
            | Bool
otherwise = LeftHandSide s v1 env1 a
-> LeftHandSide s v2 a a -> LeftHandSide s (v1, v2) env1 a
forall (s :: * -> *) v1 env env' v2 env''.
LeftHandSide s v1 env env'
-> LeftHandSide s v2 env' env''
-> LeftHandSide s (v1, v2) env env''
LeftHandSidePair LeftHandSide s v1 env1 a
l1' LeftHandSide s v2 a a
l2''
        in
          (Bool
c1 Bool -> Bool -> Bool
|| Bool
c2, Usages
us'', LeftHandSide s (v1, v2) env1 a
-> Exists (LeftHandSide s (v1, v2) env1)
forall (f :: * -> *) a. f a -> Exists f
Exists LeftHandSide s (v1, v2) env1 a
lhs')
    go Usages
_ LeftHandSide s t env1 env2
_ = String -> (Bool, Usages, Exists (LeftHandSide s t env1))
forall a. HasCallStack => String -> a
internalError String
"Empty array, mismatch in length of usages array and LHS"
shrinkLhs Count
_ LeftHandSide s t env1 env2
_ = Maybe (Exists (LeftHandSide s t env1))
forall a. Maybe a
Nothing

-- The first LHS should be 'larger' than the second, eg the second may have
-- a wildcard if the first LHS does bind variables there, but not the other
-- way around.
--
strengthenShrunkLHS
    :: HasCallStack
    => LeftHandSide s t env1 env2
    -> LeftHandSide s t env1' env2'
    -> env1 :?> env1'
    -> env2 :?> env2'
strengthenShrunkLHS :: LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS (LeftHandSideWildcard TupR s t
_) (LeftHandSideWildcard TupR s t
_) env1 :?> env1'
k = Idx env2 t' -> Maybe (Idx env2' t')
env1 :?> env1'
k
strengthenShrunkLHS (LeftHandSideSingle s t
_)   (LeftHandSideSingle s t
_)   env1 :?> env1'
k = \Idx env2 t'
ix -> case Idx env2 t'
ix of
  Idx env2 t'
ZeroIdx     -> Idx (env1', t') t' -> Maybe (Idx (env1', t') t')
forall a. a -> Maybe a
Just Idx (env1', t') t'
forall env t. Idx (env, t) t
ZeroIdx
  SuccIdx Idx env t'
ix' -> Idx env1' t' -> Idx (env1', t) t'
forall env t s. Idx env t -> Idx (env, s) t
SuccIdx (Idx env1' t' -> Idx (env1', t) t')
-> Maybe (Idx env1' t') -> Maybe (Idx (env1', t) t')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Idx env1 t' -> Maybe (Idx env1' t')
env1 :?> env1'
k Idx env1 t'
Idx env t'
ix'
strengthenShrunkLHS (LeftHandSidePair LeftHandSide s v1 env1 env'
lA LeftHandSide s v2 env' env2
hA) (LeftHandSidePair LeftHandSide s v1 env1' env'
lB LeftHandSide s v2 env' env2'
hB) env1 :?> env1'
k = LeftHandSide s v2 env' env2
-> LeftHandSide s v2 env' env2'
-> (env' :?> env')
-> env2 :?> env2'
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS LeftHandSide s v2 env' env2
hA LeftHandSide s v2 env' env2'
LeftHandSide s v2 env' env2'
hB ((env' :?> env') -> env2 :?> env2')
-> (env' :?> env') -> env2 :?> env2'
forall a b. (a -> b) -> a -> b
$ LeftHandSide s v1 env1 env'
-> LeftHandSide s v1 env1' env'
-> (env1 :?> env1')
-> env' :?> env'
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS LeftHandSide s v1 env1 env'
lA LeftHandSide s v1 env1' env'
LeftHandSide s v1 env1' env'
lB env1 :?> env1'
k
strengthenShrunkLHS (LeftHandSideSingle s t
_)   (LeftHandSideWildcard TupR s t
_) env1 :?> env1'
k = \Idx env2 t'
ix -> case Idx env2 t'
ix of
  Idx env2 t'
ZeroIdx     -> Maybe (Idx env2' t')
forall a. Maybe a
Nothing
  SuccIdx Idx env t'
ix' -> Idx env1 t' -> Maybe (Idx env1' t')
env1 :?> env1'
k Idx env1 t'
Idx env t'
ix'
strengthenShrunkLHS (LeftHandSidePair LeftHandSide s v1 env1 env'
l LeftHandSide s v2 env' env2
h)   (LeftHandSideWildcard TupR s t
t) env1 :?> env1'
k = LeftHandSide s v2 env' env2
-> LeftHandSide s v2 env2' env2'
-> (env' :?> env2')
-> env2 :?> env2'
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS LeftHandSide s v2 env' env2
h (TupR s v2 -> LeftHandSide s v2 env2' env2'
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard TupR s v2
t2) ((env' :?> env2') -> env2 :?> env2')
-> (env' :?> env2') -> env2 :?> env2'
forall a b. (a -> b) -> a -> b
$ LeftHandSide s v1 env1 env'
-> LeftHandSide s v1 env2' env2'
-> (env1 :?> env2')
-> env' :?> env2'
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS LeftHandSide s v1 env1 env'
l (TupR s v1 -> LeftHandSide s v1 env2' env2'
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard TupR s v1
t1) env1 :?> env1'
env1 :?> env2'
k
  where
    TupRpair TupR s a
t1 TupR s b
t2 = TupR s t
t
strengthenShrunkLHS (LeftHandSideWildcard TupR s t
_) LeftHandSide s t env1' env2'
_                        env1 :?> env1'
_ = String -> Idx env2 t' -> Maybe (Idx env2' t')
forall a. HasCallStack => String -> a
internalError String
"Second LHS defines more variables"
strengthenShrunkLHS LeftHandSide s t env1 env2
_                        LeftHandSide s t env1' env2'
_                        env1 :?> env1'
_ = String -> Idx env2 t' -> Maybe (Idx env2' t')
forall a. HasCallStack => String -> a
internalError String
"Mismatch LHS single with LHS pair"


-- Shrinking
-- =========

-- The shrinking substitution for scalar expressions. This is a restricted
-- instance of beta-reduction to cases where the bound variable is used zero
-- (dead-code elimination) or one (linear inlining) times.
--
shrinkExp :: HasCallStack => OpenExp env aenv t -> (Bool, OpenExp env aenv t)
shrinkExp :: OpenExp env aenv t -> (Bool, OpenExp env aenv t)
shrinkExp = Text -> (Bool, OpenExp env aenv t) -> (Bool, OpenExp env aenv t)
forall a. Text -> a -> a
Stats.substitution Text
"shrinkE" ((Bool, OpenExp env aenv t) -> (Bool, OpenExp env aenv t))
-> (OpenExp env aenv t -> (Bool, OpenExp env aenv t))
-> OpenExp env aenv t
-> (Bool, OpenExp env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Any -> Bool)
-> (Any, OpenExp env aenv t) -> (Bool, OpenExp env aenv t)
forall a a' b. (a -> a') -> (a, b) -> (a', b)
first Any -> Bool
getAny ((Any, OpenExp env aenv t) -> (Bool, OpenExp env aenv t))
-> (OpenExp env aenv t -> (Any, OpenExp env aenv t))
-> OpenExp env aenv t
-> (Bool, OpenExp env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE
  where
    -- If the bound variable is used at most this many times, it will be inlined
    -- into the body. In cases where it is not used at all, this is equivalent
    -- to dead-code elimination.
    --
    lIMIT :: Int
    lIMIT :: Int
lIMIT = Int
1

    cheap :: OpenExp env aenv t -> Bool
    cheap :: OpenExp env aenv t -> Bool
cheap (Evar ExpVar env t
_)       = Bool
True
    cheap (Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2)   = OpenExp env aenv t1 -> Bool
forall env aenv t. OpenExp env aenv t -> Bool
cheap OpenExp env aenv t1
e1 Bool -> Bool -> Bool
&& OpenExp env aenv t2 -> Bool
forall env aenv t. OpenExp env aenv t -> Bool
cheap OpenExp env aenv t2
e2
    cheap OpenExp env aenv t
Nil            = Bool
True
    cheap Const{}        = Bool
True
    cheap PrimConst{}    = Bool
True
    cheap Undef{}        = Bool
True
    cheap (Coerce ScalarType a
_ ScalarType t
_ OpenExp env aenv a
e) = OpenExp env aenv a -> Bool
forall env aenv t. OpenExp env aenv t -> Bool
cheap OpenExp env aenv a
e
    cheap OpenExp env aenv t
_              = Bool
False

    shrinkE :: HasCallStack => OpenExp env aenv t -> (Any, OpenExp env aenv t)
    shrinkE :: OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
exp = case OpenExp env aenv t
exp of
      Let (LeftHandSideSingle ScalarType bnd_t
_) bnd :: OpenExp env aenv bnd_t
bnd@Evar{} OpenExp env' aenv t
body -> Text -> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall a. Text -> a -> a
Stats.inline Text
"Var"   ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> (Any, OpenExp env aenv t)
-> (Any, OpenExp env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall x. (Any, x) -> (Any, x)
yes ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall a b. (a -> b) -> a -> b
$ OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE (OpenExp (env, bnd_t) aenv t
-> OpenExp env aenv bnd_t -> OpenExp env aenv t
forall env s aenv t.
OpenExp (env, s) aenv t -> OpenExp env aenv s -> OpenExp env aenv t
inline OpenExp env' aenv t
OpenExp (env, bnd_t) aenv t
body OpenExp env aenv bnd_t
bnd)
      Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd OpenExp env' aenv t
body
        | Bool
shouldInline -> case ELeftHandSide bnd_t env env'
-> OpenExp env' aenv t
-> OpenExp env aenv bnd_t
-> Maybe (OpenExp env aenv t)
forall env env' aenv t1 t2.
ELeftHandSide t1 env env'
-> OpenExp env' aenv t2
-> OpenExp env aenv t1
-> Maybe (OpenExp env aenv t2)
inlineVars ELeftHandSide bnd_t env env'
lhs ((Any, OpenExp env' aenv t) -> OpenExp env' aenv t
forall a b. (a, b) -> b
snd (Any, OpenExp env' aenv t)
body') ((Any, OpenExp env aenv bnd_t) -> OpenExp env aenv bnd_t
forall a b. (a, b) -> b
snd (Any, OpenExp env aenv bnd_t)
bnd') of
            Just OpenExp env aenv t
inlined -> Text -> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall a. Text -> a -> a
Stats.betaReduce Text
msg ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> (Any, OpenExp env aenv t)
-> (Any, OpenExp env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall x. (Any, x) -> (Any, x)
yes ((Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t))
-> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall a b. (a -> b) -> a -> b
$ OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
inlined
            Maybe (OpenExp env aenv t)
_            -> String -> (Any, OpenExp env aenv t)
forall a. HasCallStack => String -> a
internalError String
"Unexpected failure while trying to inline some expression."
        | Just (Exists LeftHandSide ScalarType bnd_t env a
lhs') <- Count
-> ELeftHandSide bnd_t env env'
-> Maybe (Exists (LeftHandSide ScalarType bnd_t env))
forall (s :: * -> *) t env1 env2.
HasCallStack =>
Count
-> LeftHandSide s t env1 env2
-> Maybe (Exists (LeftHandSide s t env1))
shrinkLhs Count
count ELeftHandSide bnd_t env env'
lhs -> case (env' :?> a) -> OpenExp env' aenv t -> Maybe (OpenExp a aenv t)
forall (f :: * -> * -> * -> *) env env' aenv t.
RebuildableExp f =>
(env :?> env') -> f env aenv t -> Maybe (f env' aenv t)
strengthenE (ELeftHandSide bnd_t env env'
-> LeftHandSide ScalarType bnd_t env a
-> (env :?> env)
-> env' :?> a
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS ELeftHandSide bnd_t env env'
lhs LeftHandSide ScalarType bnd_t env a
lhs' forall a. a -> Maybe a
env :?> env
Just) ((Any, OpenExp env' aenv t) -> OpenExp env' aenv t
forall a b. (a, b) -> b
snd (Any, OpenExp env' aenv t)
body') of
           Just OpenExp a aenv t
body'' -> (Bool -> Any
Any Bool
True, LeftHandSide ScalarType bnd_t env a
-> OpenExp env aenv bnd_t -> OpenExp a aenv t -> OpenExp env aenv t
forall bnd_t env env' aenv body_t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv body_t
-> OpenExp env aenv body_t
Let LeftHandSide ScalarType bnd_t env a
lhs' ((Any, OpenExp env aenv bnd_t) -> OpenExp env aenv bnd_t
forall a b. (a, b) -> b
snd (Any, OpenExp env aenv bnd_t)
bnd') OpenExp a aenv t
body'')
           Maybe (OpenExp a aenv t)
Nothing     -> String -> (Any, OpenExp env aenv t)
forall a. HasCallStack => String -> a
internalError String
"Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE."
        | Bool
otherwise    -> ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv t
-> OpenExp env aenv t
forall bnd_t env env' aenv body_t.
ELeftHandSide bnd_t env env'
-> OpenExp env aenv bnd_t
-> OpenExp env' aenv body_t
-> OpenExp env aenv body_t
Let ELeftHandSide bnd_t env env'
lhs (OpenExp env aenv bnd_t
 -> OpenExp env' aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env aenv bnd_t)
-> (Any, OpenExp env' aenv t -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Any, OpenExp env aenv bnd_t)
bnd' (Any, OpenExp env' aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env' aenv t) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Any, OpenExp env' aenv t)
body'
        where
          shouldInline :: Bool
shouldInline = case Count
count of
            Finite Int
0     -> Bool
False -- Handled by shrinkLhs
            Finite Int
n     -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
lIMIT Bool -> Bool -> Bool
|| OpenExp env aenv bnd_t -> Bool
forall env aenv t. OpenExp env aenv t -> Bool
cheap ((Any, OpenExp env aenv bnd_t) -> OpenExp env aenv bnd_t
forall a b. (a, b) -> b
snd (Any, OpenExp env aenv bnd_t)
bnd')
            Count
Infinity     ->               OpenExp env aenv bnd_t -> Bool
forall env aenv t. OpenExp env aenv t -> Bool
cheap ((Any, OpenExp env aenv bnd_t) -> OpenExp env aenv bnd_t
forall a b. (a, b) -> b
snd (Any, OpenExp env aenv bnd_t)
bnd')
            Impossible Usages
_ -> Bool
False

          bnd' :: (Any, OpenExp env aenv bnd_t)
bnd'  = OpenExp env aenv bnd_t -> (Any, OpenExp env aenv bnd_t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv bnd_t
bnd
          body' :: (Any, OpenExp env' aenv t)
body' = OpenExp env' aenv t -> (Any, OpenExp env' aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env' aenv t
body

          -- If the lhs includes non-trivial wildcards (the last field of range is Nothing),
          -- then we cannot inline the binding. We can only check which variables are not used,
          -- to detect unused variables.
          --
          -- If the lhs does not include non-trivial wildcards (the last field of range is a Just),
          -- we can both analyse whether we can inline the binding, and check which variables are
          -- not used, to detect unused variables.
          --
          count :: Count
count = case ELeftHandSide bnd_t env env'
-> Either (env :~: env') (VarsRange env')
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env')
lhsVarsRange ELeftHandSide bnd_t env env'
lhs of
            Left env :~: env'
_      -> Int -> Count
Finite Int
0
            Right VarsRange env'
range -> VarsRange env' -> OpenExp env' aenv t -> Count
forall env aenv t. VarsRange env -> OpenExp env aenv t -> Count
usesOfExp VarsRange env'
range ((Any, OpenExp env' aenv t) -> OpenExp env' aenv t
forall a b. (a, b) -> b
snd (Any, OpenExp env' aenv t)
body')

          msg :: Text
msg = case Count
count of
            Finite Int
0 -> Text
"dead exp"
            Count
_        -> Text
"inline exp"   -- forced inlining when lIMIT > 1
      --
      Evar ExpVar env t
v                    -> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpVar env t -> OpenExp env aenv t
forall env t aenv. ExpVar env t -> OpenExp env aenv t
Evar ExpVar env t
v)
      Const ScalarType t
t t
c                 -> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ScalarType t -> t -> OpenExp env aenv t
forall t env aenv. ScalarType t -> t -> OpenExp env aenv t
Const ScalarType t
t t
c)
      Undef ScalarType t
t                   -> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ScalarType t -> OpenExp env aenv t
forall t env aenv. ScalarType t -> OpenExp env aenv t
Undef ScalarType t
t)
      OpenExp env aenv t
Nil                       -> OpenExp env aenv () -> (Any, OpenExp env aenv ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure OpenExp env aenv ()
forall env aenv. OpenExp env aenv ()
Nil
      Pair OpenExp env aenv t1
x OpenExp env aenv t2
y                  -> OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
forall env aenv t1 t2.
OpenExp env aenv t1
-> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2)
Pair (OpenExp env aenv t1
 -> OpenExp env aenv t2 -> OpenExp env aenv (t1, t2))
-> (Any, OpenExp env aenv t1)
-> (Any, OpenExp env aenv t2 -> OpenExp env aenv (t1, t2))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv t1 -> (Any, OpenExp env aenv t1)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t1
x (Any, OpenExp env aenv t2 -> OpenExp env aenv (t1, t2))
-> (Any, OpenExp env aenv t2) -> (Any, OpenExp env aenv (t1, t2))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv t2 -> (Any, OpenExp env aenv t2)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t2
y
      VecPack   VecR n s tup
vec OpenExp env aenv tup
e           -> VecR n s tup -> OpenExp env aenv tup -> OpenExp env aenv (Vec n s)
forall (n :: Nat) s tup env aenv.
KnownNat n =>
VecR n s tup -> OpenExp env aenv tup -> OpenExp env aenv (Vec n s)
VecPack   VecR n s tup
vec (OpenExp env aenv tup -> OpenExp env aenv (Vec n s))
-> (Any, OpenExp env aenv tup) -> (Any, OpenExp env aenv (Vec n s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv tup -> (Any, OpenExp env aenv tup)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv tup
e
      VecUnpack VecR n s t
vec OpenExp env aenv (Vec n s)
e           -> VecR n s t -> OpenExp env aenv (Vec n s) -> OpenExp env aenv t
forall (n :: Nat) s tup env aenv.
KnownNat n =>
VecR n s tup -> OpenExp env aenv (Vec n s) -> OpenExp env aenv tup
VecUnpack VecR n s t
vec (OpenExp env aenv (Vec n s) -> OpenExp env aenv t)
-> (Any, OpenExp env aenv (Vec n s)) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv (Vec n s) -> (Any, OpenExp env aenv (Vec n s))
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv (Vec n s)
e
      IndexSlice SliceIndex slix t co sh
x OpenExp env aenv slix
ix OpenExp env aenv sh
sh        -> SliceIndex slix t co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv t
forall slix sl co sh env aenv.
SliceIndex slix sl co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sh
-> OpenExp env aenv sl
IndexSlice SliceIndex slix t co sh
x (OpenExp env aenv slix
 -> OpenExp env aenv sh -> OpenExp env aenv t)
-> (Any, OpenExp env aenv slix)
-> (Any, OpenExp env aenv sh -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv slix -> (Any, OpenExp env aenv slix)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv slix
ix (Any, OpenExp env aenv sh -> OpenExp env aenv t)
-> (Any, OpenExp env aenv sh) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv sh -> (Any, OpenExp env aenv sh)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv sh
sh
      IndexFull SliceIndex slix sl co t
x OpenExp env aenv slix
ix OpenExp env aenv sl
sl         -> SliceIndex slix sl co t
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv t
forall slix sl co sh env aenv.
SliceIndex slix sl co sh
-> OpenExp env aenv slix
-> OpenExp env aenv sl
-> OpenExp env aenv sh
IndexFull SliceIndex slix sl co t
x (OpenExp env aenv slix
 -> OpenExp env aenv sl -> OpenExp env aenv t)
-> (Any, OpenExp env aenv slix)
-> (Any, OpenExp env aenv sl -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv slix -> (Any, OpenExp env aenv slix)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv slix
ix (Any, OpenExp env aenv sl -> OpenExp env aenv t)
-> (Any, OpenExp env aenv sl) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv sl -> (Any, OpenExp env aenv sl)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv sl
sl
      ToIndex ShapeR sh
shr OpenExp env aenv sh
sh OpenExp env aenv sh
ix         -> ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
ToIndex ShapeR sh
shr (OpenExp env aenv sh
 -> OpenExp env aenv sh -> OpenExp env aenv Int)
-> (Any, OpenExp env aenv sh)
-> (Any, OpenExp env aenv sh -> OpenExp env aenv Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv sh -> (Any, OpenExp env aenv sh)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv sh
sh (Any, OpenExp env aenv sh -> OpenExp env aenv Int)
-> (Any, OpenExp env aenv sh) -> (Any, OpenExp env aenv Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv sh -> (Any, OpenExp env aenv sh)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv sh
ix
      FromIndex ShapeR t
shr OpenExp env aenv t
sh OpenExp env aenv Int
i        -> ShapeR t
-> OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t
forall sh env aenv.
ShapeR sh
-> OpenExp env aenv sh
-> OpenExp env aenv Int
-> OpenExp env aenv sh
FromIndex ShapeR t
shr (OpenExp env aenv t -> OpenExp env aenv Int -> OpenExp env aenv t)
-> (Any, OpenExp env aenv t)
-> (Any, OpenExp env aenv Int -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
sh (Any, OpenExp env aenv Int -> OpenExp env aenv t)
-> (Any, OpenExp env aenv Int) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv Int -> (Any, OpenExp env aenv Int)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv Int
i
      Case OpenExp env aenv TAG
e [(TAG, OpenExp env aenv t)]
rhs Maybe (OpenExp env aenv t)
def            -> OpenExp env aenv TAG
-> [(TAG, OpenExp env aenv t)]
-> Maybe (OpenExp env aenv t)
-> OpenExp env aenv t
forall env aenv b.
OpenExp env aenv TAG
-> [(TAG, OpenExp env aenv b)]
-> Maybe (OpenExp env aenv b)
-> OpenExp env aenv b
Case (OpenExp env aenv TAG
 -> [(TAG, OpenExp env aenv t)]
 -> Maybe (OpenExp env aenv t)
 -> OpenExp env aenv t)
-> (Any, OpenExp env aenv TAG)
-> (Any,
    [(TAG, OpenExp env aenv t)]
    -> Maybe (OpenExp env aenv t) -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv TAG -> (Any, OpenExp env aenv TAG)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv TAG
e (Any,
 [(TAG, OpenExp env aenv t)]
 -> Maybe (OpenExp env aenv t) -> OpenExp env aenv t)
-> (Any, [(TAG, OpenExp env aenv t)])
-> (Any, Maybe (OpenExp env aenv t) -> OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(Any, (TAG, OpenExp env aenv t))]
-> (Any, [(TAG, OpenExp env aenv t)])
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA [ (TAG
t,) (OpenExp env aenv t -> (TAG, OpenExp env aenv t))
-> (Any, OpenExp env aenv t) -> (Any, (TAG, OpenExp env aenv t))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
c | (TAG
t,OpenExp env aenv t
c) <- [(TAG, OpenExp env aenv t)]
rhs ] (Any, Maybe (OpenExp env aenv t) -> OpenExp env aenv t)
-> (Any, Maybe (OpenExp env aenv t)) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
forall env aenv t.
HasCallStack =>
Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
shrinkMaybeE Maybe (OpenExp env aenv t)
def
      Cond OpenExp env aenv TAG
p OpenExp env aenv t
t OpenExp env aenv t
e                -> OpenExp env aenv TAG
-> OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
forall env aenv t.
OpenExp env aenv TAG
-> OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t
Cond (OpenExp env aenv TAG
 -> OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env aenv TAG)
-> (Any,
    OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv TAG -> (Any, OpenExp env aenv TAG)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv TAG
p (Any,
 OpenExp env aenv t -> OpenExp env aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env aenv t)
-> (Any, OpenExp env aenv t -> OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
t (Any, OpenExp env aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
e
      While OpenFun env aenv (t -> TAG)
p OpenFun env aenv (t -> t)
f OpenExp env aenv t
x               -> OpenFun env aenv (t -> TAG)
-> OpenFun env aenv (t -> t)
-> OpenExp env aenv t
-> OpenExp env aenv t
forall env aenv a.
OpenFun env aenv (a -> TAG)
-> OpenFun env aenv (a -> a)
-> OpenExp env aenv a
-> OpenExp env aenv a
While (OpenFun env aenv (t -> TAG)
 -> OpenFun env aenv (t -> t)
 -> OpenExp env aenv t
 -> OpenExp env aenv t)
-> (Any, OpenFun env aenv (t -> TAG))
-> (Any,
    OpenFun env aenv (t -> t)
    -> OpenExp env aenv t -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenFun env aenv (t -> TAG) -> (Any, OpenFun env aenv (t -> TAG))
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> (Any, OpenFun env aenv t)
shrinkF OpenFun env aenv (t -> TAG)
p (Any,
 OpenFun env aenv (t -> t)
 -> OpenExp env aenv t -> OpenExp env aenv t)
-> (Any, OpenFun env aenv (t -> t))
-> (Any, OpenExp env aenv t -> OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenFun env aenv (t -> t) -> (Any, OpenFun env aenv (t -> t))
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> (Any, OpenFun env aenv t)
shrinkF OpenFun env aenv (t -> t)
f (Any, OpenExp env aenv t -> OpenExp env aenv t)
-> (Any, OpenExp env aenv t) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
x
      PrimConst PrimConst t
c               -> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimConst t -> OpenExp env aenv t
forall t env aenv. PrimConst t -> OpenExp env aenv t
PrimConst PrimConst t
c)
      PrimApp PrimFun (a -> t)
f OpenExp env aenv a
x               -> PrimFun (a -> t) -> OpenExp env aenv a -> OpenExp env aenv t
forall a r env aenv.
PrimFun (a -> r) -> OpenExp env aenv a -> OpenExp env aenv r
PrimApp PrimFun (a -> t)
f (OpenExp env aenv a -> OpenExp env aenv t)
-> (Any, OpenExp env aenv a) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv a -> (Any, OpenExp env aenv a)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv a
x
      Index ArrayVar aenv (Array dim t)
a OpenExp env aenv dim
sh                -> ArrayVar aenv (Array dim t)
-> OpenExp env aenv dim -> OpenExp env aenv t
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv dim -> OpenExp env aenv t
Index ArrayVar aenv (Array dim t)
a (OpenExp env aenv dim -> OpenExp env aenv t)
-> (Any, OpenExp env aenv dim) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv dim -> (Any, OpenExp env aenv dim)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv dim
sh
      LinearIndex ArrayVar aenv (Array dim t)
a OpenExp env aenv Int
i           -> ArrayVar aenv (Array dim t)
-> OpenExp env aenv Int -> OpenExp env aenv t
forall aenv dim t env.
ArrayVar aenv (Array dim t)
-> OpenExp env aenv Int -> OpenExp env aenv t
LinearIndex ArrayVar aenv (Array dim t)
a (OpenExp env aenv Int -> OpenExp env aenv t)
-> (Any, OpenExp env aenv Int) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv Int -> (Any, OpenExp env aenv Int)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv Int
i
      Shape ArrayVar aenv (Array t e)
a                   -> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ArrayVar aenv (Array t e) -> OpenExp env aenv t
forall aenv dim e env.
ArrayVar aenv (Array dim e) -> OpenExp env aenv dim
Shape ArrayVar aenv (Array t e)
a)
      ShapeSize ShapeR dim
shr OpenExp env aenv dim
sh          -> ShapeR dim -> OpenExp env aenv dim -> OpenExp env aenv Int
forall dim env aenv.
ShapeR dim -> OpenExp env aenv dim -> OpenExp env aenv Int
ShapeSize ShapeR dim
shr (OpenExp env aenv dim -> OpenExp env aenv Int)
-> (Any, OpenExp env aenv dim) -> (Any, OpenExp env aenv Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv dim -> (Any, OpenExp env aenv dim)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv dim
sh
      Foreign TypeR t
repr asm (x -> t)
ff Fun () (x -> t)
f OpenExp env aenv x
e       -> TypeR t
-> asm (x -> t)
-> Fun () (x -> t)
-> OpenExp env aenv x
-> OpenExp env aenv t
forall (asm :: * -> *) y x env aenv.
Foreign asm =>
TypeR y
-> asm (x -> y)
-> Fun () (x -> y)
-> OpenExp env aenv x
-> OpenExp env aenv y
Foreign TypeR t
repr asm (x -> t)
ff (Fun () (x -> t) -> OpenExp env aenv x -> OpenExp env aenv t)
-> (Any, Fun () (x -> t))
-> (Any, OpenExp env aenv x -> OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fun () (x -> t) -> (Any, Fun () (x -> t))
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> (Any, OpenFun env aenv t)
shrinkF Fun () (x -> t)
f (Any, OpenExp env aenv x -> OpenExp env aenv t)
-> (Any, OpenExp env aenv x) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OpenExp env aenv x -> (Any, OpenExp env aenv x)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv x
e
      Coerce ScalarType a
t1 ScalarType t
t2 OpenExp env aenv a
e            -> ScalarType a
-> ScalarType t -> OpenExp env aenv a -> OpenExp env aenv t
forall a b env aenv.
BitSizeEq a b =>
ScalarType a
-> ScalarType b -> OpenExp env aenv a -> OpenExp env aenv b
Coerce ScalarType a
t1 ScalarType t
t2 (OpenExp env aenv a -> OpenExp env aenv t)
-> (Any, OpenExp env aenv a) -> (Any, OpenExp env aenv t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv a -> (Any, OpenExp env aenv a)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv a
e

    shrinkF :: HasCallStack => OpenFun env aenv t -> (Any, OpenFun env aenv t)
    shrinkF :: OpenFun env aenv t -> (Any, OpenFun env aenv t)
shrinkF = (Bool -> Any)
-> (Bool, OpenFun env aenv t) -> (Any, OpenFun env aenv t)
forall a a' b. (a -> a') -> (a, b) -> (a', b)
first Bool -> Any
Any ((Bool, OpenFun env aenv t) -> (Any, OpenFun env aenv t))
-> (OpenFun env aenv t -> (Bool, OpenFun env aenv t))
-> OpenFun env aenv t
-> (Any, OpenFun env aenv t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenFun env aenv t -> (Bool, OpenFun env aenv t)
forall env aenv f.
HasCallStack =>
OpenFun env aenv f -> (Bool, OpenFun env aenv f)
shrinkFun

    shrinkMaybeE :: HasCallStack => Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
    shrinkMaybeE :: Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
shrinkMaybeE Maybe (OpenExp env aenv t)
Nothing  = Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (OpenExp env aenv t)
forall a. Maybe a
Nothing
    shrinkMaybeE (Just OpenExp env aenv t
e) = OpenExp env aenv t -> Maybe (OpenExp env aenv t)
forall a. a -> Maybe a
Just (OpenExp env aenv t -> Maybe (OpenExp env aenv t))
-> (Any, OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv t -> (Any, OpenExp env aenv t)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Any, OpenExp env aenv t)
shrinkE OpenExp env aenv t
e

    first :: (a -> a') -> (a,b) -> (a',b)
    first :: (a -> a') -> (a, b) -> (a', b)
first a -> a'
f (a
x,b
y) = (a -> a'
f a
x, b
y)

    yes :: (Any, x) -> (Any, x)
    yes :: (Any, x) -> (Any, x)
yes (Any
_, x
x) = (Bool -> Any
Any Bool
True, x
x)

shrinkFun :: HasCallStack => OpenFun env aenv f -> (Bool, OpenFun env aenv f)
shrinkFun :: OpenFun env aenv f -> (Bool, OpenFun env aenv f)
shrinkFun (Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t
f) = case ELeftHandSide a env env' -> Either (env :~: env') (VarsRange env')
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env')
lhsVarsRange ELeftHandSide a env env'
lhs of
  Left env :~: env'
Refl ->
    let b' :: Bool
b' = case ELeftHandSide a env env'
lhs of
                LeftHandSideWildcard TupR ScalarType a
TupRunit -> Bool
b
                ELeftHandSide a env env'
_                             -> Bool
True
    in (Bool
b', ELeftHandSide a env' env'
-> OpenFun env' aenv t -> OpenFun env' aenv (a -> t)
forall a env env' aenv t.
ELeftHandSide a env env'
-> OpenFun env' aenv t -> OpenFun env aenv (a -> t)
Lam (TupR ScalarType a -> ELeftHandSide a env' env'
forall (s :: * -> *) v env. TupR s v -> LeftHandSide s v env env
LeftHandSideWildcard (TupR ScalarType a -> ELeftHandSide a env' env')
-> TupR ScalarType a -> ELeftHandSide a env' env'
forall a b. (a -> b) -> a -> b
$ ELeftHandSide a env env' -> TupR ScalarType a
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> TupR s v
lhsToTupR ELeftHandSide a env env'
lhs) OpenFun env' aenv t
f')
  Right VarsRange env'
range ->
    let
      count :: Count
count = VarsRange env' -> OpenFun env' aenv t -> Count
forall env aenv f. VarsRange env -> OpenFun env aenv f -> Count
usesOfFun VarsRange env'
range OpenFun env' aenv t
f
    in case Count
-> ELeftHandSide a env env'
-> Maybe (Exists (LeftHandSide ScalarType a env))
forall (s :: * -> *) t env1 env2.
HasCallStack =>
Count
-> LeftHandSide s t env1 env2
-> Maybe (Exists (LeftHandSide s t env1))
shrinkLhs Count
count ELeftHandSide a env env'
lhs of
        Just (Exists LeftHandSide ScalarType a env a
lhs') -> case (env' :?> a) -> OpenFun env' aenv t -> Maybe (OpenFun a aenv t)
forall (f :: * -> * -> * -> *) env env' aenv t.
RebuildableExp f =>
(env :?> env') -> f env aenv t -> Maybe (f env' aenv t)
strengthenE (ELeftHandSide a env env'
-> LeftHandSide ScalarType a env a -> (env :?> env) -> env' :?> a
forall (s :: * -> *) t env1 env2 env1' env2'.
HasCallStack =>
LeftHandSide s t env1 env2
-> LeftHandSide s t env1' env2'
-> (env1 :?> env1')
-> env2 :?> env2'
strengthenShrunkLHS ELeftHandSide a env env'
lhs LeftHandSide ScalarType a env a
lhs' forall a. a -> Maybe a
env :?> env
Just) OpenFun env' aenv t
f' of
          Just OpenFun a aenv t
f'' -> (Bool
True, LeftHandSide ScalarType a env a
-> OpenFun a aenv t -> OpenFun env aenv (a -> t)
forall a env env' aenv t.
ELeftHandSide a env env'
-> OpenFun env' aenv t -> OpenFun env aenv (a -> t)
Lam LeftHandSide ScalarType a env a
lhs' OpenFun a aenv t
f'')
          Maybe (OpenFun a aenv t)
Nothing  -> String -> (Bool, OpenFun env aenv f)
forall a. HasCallStack => String -> a
internalError String
"Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE."
        Maybe (Exists (LeftHandSide ScalarType a env))
Nothing -> (Bool
b, ELeftHandSide a env env'
-> OpenFun env' aenv t -> OpenFun env aenv (a -> t)
forall a env env' aenv t.
ELeftHandSide a env env'
-> OpenFun env' aenv t -> OpenFun env aenv (a -> t)
Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t
f')
  where
    (Bool
b, OpenFun env' aenv t
f') = OpenFun env' aenv t -> (Bool, OpenFun env' aenv t)
forall env aenv f.
HasCallStack =>
OpenFun env aenv f -> (Bool, OpenFun env aenv f)
shrinkFun OpenFun env' aenv t
f

shrinkFun (Body OpenExp env aenv f
b) = OpenExp env aenv f -> OpenFun env aenv f
forall env aenv t. OpenExp env aenv t -> OpenFun env aenv t
Body (OpenExp env aenv f -> OpenFun env aenv f)
-> (Bool, OpenExp env aenv f) -> (Bool, OpenFun env aenv f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenExp env aenv f -> (Bool, OpenExp env aenv f)
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> (Bool, OpenExp env aenv t)
shrinkExp OpenExp env aenv f
b

-- The shrinking substitution for array computations. This is further limited to
-- dead-code elimination only, primarily because linear inlining may inline
-- array computations into scalar expressions, which is generally not desirable.
--
type ShrinkAcc acc = forall aenv a. acc aenv a -> acc aenv a

{--
type ReduceAcc acc = forall aenv s t. acc aenv s -> acc (aenv,s) t -> Maybe (PreOpenAcc acc aenv t)

shrinkPreAcc
    :: forall acc aenv arrs. ShrinkAcc acc -> ReduceAcc acc
    -> PreOpenAcc acc aenv arrs
    -> PreOpenAcc acc aenv arrs
shrinkPreAcc shrinkAcc reduceAcc = Stats.substitution "shrinkA" shrinkA
  where
    shrinkA :: PreOpenAcc acc aenv' a -> PreOpenAcc acc aenv' a
    shrinkA pacc = case pacc of
      Alet lhs bnd body
        | Just reduct <- reduceAcc bnd' body'   -> shrinkA reduct
        | otherwise                             -> Alet lhs bnd' body'
        where
          bnd'  = shrinkAcc bnd
          body' = shrinkAcc body
      --
      Avar ix                   -> Avar ix
      Apair a1 a2               -> Apair (shrinkAcc a1) (shrinkAcc a2)
      Anil                      -> Anil
      Apply repr f a            -> Apply repr (shrinkAF f) (shrinkAcc a)
      Aforeign ff af a          -> Aforeign ff af (shrinkAcc a)
      Acond p t e               -> Acond (shrinkE p) (shrinkAcc t) (shrinkAcc e)
      Awhile p f a              -> Awhile (shrinkAF p) (shrinkAF f) (shrinkAcc a)
      Use repr a                -> Use repr a
      Unit e                    -> Unit (shrinkE e)
      Reshape e a               -> Reshape (shrinkE e) (shrinkAcc a)
      Generate e f              -> Generate (shrinkE e) (shrinkF f)
      Transform sh ix f a       -> Transform (shrinkE sh) (shrinkF ix) (shrinkF f) (shrinkAcc a)
      Replicate sl slix a       -> Replicate sl (shrinkE slix) (shrinkAcc a)
      Slice sl a slix           -> Slice sl (shrinkAcc a) (shrinkE slix)
      Map f a                   -> Map (shrinkF f) (shrinkAcc a)
      ZipWith f a1 a2           -> ZipWith (shrinkF f) (shrinkAcc a1) (shrinkAcc a2)
      Fold f z a                -> Fold (shrinkF f) (shrinkE z) (shrinkAcc a)
      Fold1 f a                 -> Fold1 (shrinkF f) (shrinkAcc a)
      FoldSeg f z a b           -> FoldSeg (shrinkF f) (shrinkE z) (shrinkAcc a) (shrinkAcc b)
      Fold1Seg f a b            -> Fold1Seg (shrinkF f) (shrinkAcc a) (shrinkAcc b)
      Scanl f z a               -> Scanl (shrinkF f) (shrinkE z) (shrinkAcc a)
      Scanl' f z a              -> Scanl' (shrinkF f) (shrinkE z) (shrinkAcc a)
      Scanl1 f a                -> Scanl1 (shrinkF f) (shrinkAcc a)
      Scanr f z a               -> Scanr (shrinkF f) (shrinkE z) (shrinkAcc a)
      Scanr' f z a              -> Scanr' (shrinkF f) (shrinkE z) (shrinkAcc a)
      Scanr1 f a                -> Scanr1 (shrinkF f) (shrinkAcc a)
      Permute f1 a1 f2 a2       -> Permute (shrinkF f1) (shrinkAcc a1) (shrinkF f2) (shrinkAcc a2)
      Backpermute sh f a        -> Backpermute (shrinkE sh) (shrinkF f) (shrinkAcc a)
      Stencil f b a             -> Stencil (shrinkF f) b (shrinkAcc a)
      Stencil2 f b1 a1 b2 a2    -> Stencil2 (shrinkF f) b1 (shrinkAcc a1) b2 (shrinkAcc a2)
      -- Collect s                 -> Collect (shrinkS s)

{--
    shrinkS :: PreOpenSeq acc aenv' senv a -> PreOpenSeq acc aenv' senv a
    shrinkS seq =
      case seq of
        Producer p s -> Producer (shrinkP p) (shrinkS s)
        Consumer c   -> Consumer (shrinkC c)
        Reify ix     -> Reify ix

    shrinkP :: Producer acc aenv' senv a -> Producer acc aenv' senv a
    shrinkP p =
      case p of
        StreamIn arrs        -> StreamIn arrs
        ToSeq sl slix a      -> ToSeq sl slix (shrinkAcc a)
        MapSeq f x           -> MapSeq (shrinkAF f) x
        ChunkedMapSeq f x    -> ChunkedMapSeq (shrinkAF f) x
        ZipWithSeq f x y     -> ZipWithSeq (shrinkAF f) x y
        ScanSeq f e x        -> ScanSeq (shrinkF f) (shrinkE e) x

    shrinkC :: Consumer acc aenv' senv a -> Consumer acc aenv' senv a
    shrinkC c =
      case c of
        FoldSeq f e x        -> FoldSeq (shrinkF f) (shrinkE e) x
        FoldSeqFlatten f a x -> FoldSeqFlatten (shrinkAF f) (shrinkAcc a) x
        Stuple t             -> Stuple (shrinkCT t)

    shrinkCT :: Atuple (Consumer acc aenv' senv) t -> Atuple (Consumer acc aenv' senv) t
    shrinkCT NilAtup        = NilAtup
    shrinkCT (SnocAtup t c) = SnocAtup (shrinkCT t) (shrinkC c)
--}

    shrinkE :: OpenExp env aenv' t -> OpenExp env aenv' t
    shrinkE exp = case exp of
      Let bnd body              -> Let (shrinkE bnd) (shrinkE body)
      Var idx                   -> Var idx
      Const c                   -> Const c
      Undef                     -> Undef
      Tuple t                   -> Tuple (shrinkT t)
      Prj tup e                 -> Prj tup (shrinkE e)
      IndexNil                  -> IndexNil
      IndexCons sl sz           -> IndexCons (shrinkE sl) (shrinkE sz)
      IndexHead sh              -> IndexHead (shrinkE sh)
      IndexTail sh              -> IndexTail (shrinkE sh)
      IndexSlice x ix sh        -> IndexSlice x (shrinkE ix) (shrinkE sh)
      IndexFull x ix sl         -> IndexFull x (shrinkE ix) (shrinkE sl)
      IndexAny                  -> IndexAny
      ToIndex sh ix             -> ToIndex (shrinkE sh) (shrinkE ix)
      FromIndex sh i            -> FromIndex (shrinkE sh) (shrinkE i)
      Cond p t e                -> Cond (shrinkE p) (shrinkE t) (shrinkE e)
      While p f x               -> While (shrinkF p) (shrinkF f) (shrinkE x)
      PrimConst c               -> PrimConst c
      PrimApp f x               -> PrimApp f (shrinkE x)
      Index a sh                -> Index (shrinkAcc a) (shrinkE sh)
      LinearIndex a i           -> LinearIndex (shrinkAcc a) (shrinkE i)
      Shape a                   -> Shape (shrinkAcc a)
      ShapeSize sh              -> ShapeSize (shrinkE sh)
      Intersect sh sz           -> Intersect (shrinkE sh) (shrinkE sz)
      Union sh sz               -> Union (shrinkE sh) (shrinkE sz)
      Foreign ff f e            -> Foreign ff (shrinkF f) (shrinkE e)
      Coerce e                  -> Coerce (shrinkE e)

    shrinkF :: OpenFun env aenv' f -> OpenFun env aenv' f
    shrinkF (Lam f)  = Lam (shrinkF f)
    shrinkF (Body b) = Body (shrinkE b)

    shrinkT :: Tuple (OpenExp env aenv') t -> Tuple (OpenExp env aenv') t
    shrinkT NilTup        = NilTup
    shrinkT (SnocTup t e) = shrinkT t `SnocTup` shrinkE e

    shrinkAF :: PreOpenAfun acc aenv' f -> PreOpenAfun acc aenv' f
    shrinkAF (Alam lhs f) = Alam lhs (shrinkAF f)
    shrinkAF (Abody a) = Abody (shrinkAcc a)
--}

-- Occurrence Counting
-- ===================

-- Count the number of occurrences an in-scope scalar expression bound at the
-- given variable index recursively in a term.
--
usesOfExp :: forall env aenv t. VarsRange env -> OpenExp env aenv t -> Count
usesOfExp :: VarsRange env -> OpenExp env aenv t -> Count
usesOfExp VarsRange env
range = OpenExp env aenv t -> Count
forall e. OpenExp env aenv e -> Count
countE
  where
    countE :: OpenExp env aenv e -> Count
    countE :: OpenExp env aenv e -> Count
countE OpenExp env aenv e
exp | VarsRange env -> OpenExp env aenv e -> Bool
forall env aenv t. VarsRange env -> OpenExp env aenv t -> Bool
matchEVarsRange VarsRange env
range OpenExp env aenv e
exp = Int -> Count
Finite Int
1
    countE OpenExp env aenv e
exp = case OpenExp env aenv e
exp of
      Evar ExpVar env e
v -> case VarsRange env -> ExpVar env e -> Maybe Usages
forall env (s :: * -> *) t.
VarsRange env -> Var s env t -> Maybe Usages
varInRange VarsRange env
range ExpVar env e
v of
        Just Usages
cs                 -> Usages -> Count
Impossible Usages
cs
        Maybe Usages
Nothing                 -> Int -> Count
Finite Int
0
      --
      Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
bnd OpenExp env' aenv e
body          -> OpenExp env aenv bnd_t -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv bnd_t
bnd Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> VarsRange env' -> OpenExp env' aenv e -> Count
forall env aenv t. VarsRange env -> OpenExp env aenv t -> Count
usesOfExp (ELeftHandSide bnd_t env env' -> VarsRange env -> VarsRange env'
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> VarsRange env -> VarsRange env'
weakenVarsRange ELeftHandSide bnd_t env env'
lhs VarsRange env
range) OpenExp env' aenv e
body
      Const ScalarType e
_ e
_                 -> Int -> Count
Finite Int
0
      Undef ScalarType e
_                   -> Int -> Count
Finite Int
0
      OpenExp env aenv e
Nil                       -> Int -> Count
Finite Int
0
      Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2                -> OpenExp env aenv t1 -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv t1
e1 Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv t2 -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv t2
e2
      VecPack   VecR n s tup
_ OpenExp env aenv tup
e             -> OpenExp env aenv tup -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv tup
e
      VecUnpack VecR n s e
_ OpenExp env aenv (Vec n s)
e             -> OpenExp env aenv (Vec n s) -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv (Vec n s)
e
      IndexSlice SliceIndex slix e co sh
_ OpenExp env aenv slix
ix OpenExp env aenv sh
sh        -> OpenExp env aenv slix -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv slix
ix Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv sh -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv sh
sh
      IndexFull SliceIndex slix sl co e
_ OpenExp env aenv slix
ix OpenExp env aenv sl
sl         -> OpenExp env aenv slix -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv slix
ix Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv sl -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv sl
sl
      FromIndex ShapeR e
_ OpenExp env aenv e
sh OpenExp env aenv Int
i          -> OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv e
sh Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv Int -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv Int
i
      ToIndex ShapeR sh
_ OpenExp env aenv sh
sh OpenExp env aenv sh
e            -> OpenExp env aenv sh -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv sh
sh Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv sh -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv sh
e
      Case OpenExp env aenv TAG
e [(TAG, OpenExp env aenv e)]
rhs Maybe (OpenExp env aenv e)
def            -> OpenExp env aenv TAG -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv TAG
e  Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> [Count] -> Count
forall a. Monoid a => [a] -> a
mconcat [ OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv e
c | (TAG
_,OpenExp env aenv e
c) <- [(TAG, OpenExp env aenv e)]
rhs ] Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> Count
-> (OpenExp env aenv e -> Count)
-> Maybe (OpenExp env aenv e)
-> Count
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> Count
Finite Int
0) OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE Maybe (OpenExp env aenv e)
def
      Cond OpenExp env aenv TAG
p OpenExp env aenv e
t OpenExp env aenv e
e                -> OpenExp env aenv TAG -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv TAG
p  Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv e
t Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv e
e
      While OpenFun env aenv (e -> TAG)
p OpenFun env aenv (e -> e)
f OpenExp env aenv e
x               -> OpenExp env aenv e -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv e
x  Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> Count -> Count
loopCount (VarsRange env -> OpenFun env aenv (e -> TAG) -> Count
forall env aenv f. VarsRange env -> OpenFun env aenv f -> Count
usesOfFun VarsRange env
range OpenFun env aenv (e -> TAG)
p) Count -> Count -> Count
forall a. Semigroup a => a -> a -> a
<> Count -> Count
loopCount (VarsRange env -> OpenFun env aenv (e -> e) -> Count
forall env aenv f. VarsRange env -> OpenFun env aenv f -> Count
usesOfFun VarsRange env
range OpenFun env aenv (e -> e)
f)
      PrimConst PrimConst e
_               -> Int -> Count
Finite Int
0
      PrimApp PrimFun (a -> e)
_ OpenExp env aenv a
x               -> OpenExp env aenv a -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv a
x
      Index ArrayVar aenv (Array dim e)
_ OpenExp env aenv dim
sh                -> OpenExp env aenv dim -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv dim
sh
      LinearIndex ArrayVar aenv (Array dim e)
_ OpenExp env aenv Int
i           -> OpenExp env aenv Int -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv Int
i
      Shape ArrayVar aenv (Array e e)
_                   -> Int -> Count
Finite Int
0
      ShapeSize ShapeR dim
_ OpenExp env aenv dim
sh            -> OpenExp env aenv dim -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv dim
sh
      Foreign TypeR e
_ asm (x -> e)
_ Fun () (x -> e)
_ OpenExp env aenv x
e           -> OpenExp env aenv x -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv x
e
      Coerce ScalarType a
_ ScalarType e
_ OpenExp env aenv a
e              -> OpenExp env aenv a -> Count
forall e. OpenExp env aenv e -> Count
countE OpenExp env aenv a
e

usesOfFun :: VarsRange env -> OpenFun env aenv f -> Count
usesOfFun :: VarsRange env -> OpenFun env aenv f -> Count
usesOfFun VarsRange env
range (Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t
f) = VarsRange env' -> OpenFun env' aenv t -> Count
forall env aenv f. VarsRange env -> OpenFun env aenv f -> Count
usesOfFun (ELeftHandSide a env env' -> VarsRange env -> VarsRange env'
forall (s :: * -> *) v env env'.
LeftHandSide s v env env' -> VarsRange env -> VarsRange env'
weakenVarsRange ELeftHandSide a env env'
lhs VarsRange env
range) OpenFun env' aenv t
f
usesOfFun VarsRange env
range (Body OpenExp env aenv f
b)    = VarsRange env -> OpenExp env aenv f -> Count
forall env aenv t. VarsRange env -> OpenExp env aenv t -> Count
usesOfExp VarsRange env
range OpenExp env aenv f
b

-- Count the number of occurrences of the array term bound at the given
-- environment index. If the first argument is 'True' then it includes in the
-- total uses of the variable for 'Shape' information, otherwise not.
--
type UsesOfAcc acc = forall aenv s t. Bool -> Idx aenv s -> acc aenv t -> Int

-- XXX: Should this be converted to use the above 'Count' semigroup?
--
usesOfPreAcc
    :: forall acc aenv s t.
       Bool
    -> UsesOfAcc  acc
    -> Idx            aenv s
    -> PreOpenAcc acc aenv t
    -> Int
usesOfPreAcc :: Bool -> UsesOfAcc acc -> Idx aenv s -> PreOpenAcc acc aenv t -> Int
usesOfPreAcc Bool
withShape UsesOfAcc acc
countAcc Idx aenv s
idx = PreOpenAcc acc aenv t -> Int
forall a. PreOpenAcc acc aenv a -> Int
count
  where
    countIdx :: Idx aenv a -> Int
    countIdx :: Idx aenv a -> Int
countIdx Idx aenv a
this
        | Just a :~: s
Refl <- Idx aenv a -> Idx aenv s -> Maybe (a :~: s)
forall env s t. Idx env s -> Idx env t -> Maybe (s :~: t)
matchIdx Idx aenv a
this Idx aenv s
idx = Int
1
        | Bool
otherwise                      = Int
0

    count :: PreOpenAcc acc aenv a -> Int
    count :: PreOpenAcc acc aenv a -> Int
count PreOpenAcc acc aenv a
pacc = case PreOpenAcc acc aenv a
pacc of
      Avar ArrayVar aenv (Array sh e)
var                   -> ArrayVar aenv (Array sh e) -> Int
forall a. ArrayVar aenv a -> Int
countAvar ArrayVar aenv (Array sh e)
var
      --
      Alet ALeftHandSide bndArrs aenv aenv'
lhs acc aenv bndArrs
bnd acc aenv' a
body          -> acc aenv bndArrs -> Int
forall a. acc aenv a -> Int
countA acc aenv bndArrs
bnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Bool -> Idx aenv' s -> acc aenv' a -> Int
UsesOfAcc acc
countAcc Bool
withShape (ALeftHandSide bndArrs aenv aenv' -> aenv :> aenv'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ALeftHandSide bndArrs aenv aenv'
lhs (aenv :> aenv') -> Idx aenv s -> Idx aenv' s
forall env env'.
(env :> env') -> forall t'. Idx env t' -> Idx env' t'
>:> Idx aenv s
idx) acc aenv' a
body
      Apair acc aenv as
a1 acc aenv bs
a2                -> acc aenv as -> Int
forall a. acc aenv a -> Int
countA acc aenv as
a1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv bs -> Int
forall a. acc aenv a -> Int
countA acc aenv bs
a2
      PreOpenAcc acc aenv a
Anil                       -> Int
0
      Apply ArraysR a
_ PreOpenAfun acc aenv (arrs1 -> a)
f acc aenv arrs1
a                -> PreOpenAfun acc aenv (arrs1 -> a) -> Idx aenv s -> Int
forall aenv' f. PreOpenAfun acc aenv' f -> Idx aenv' s -> Int
countAF PreOpenAfun acc aenv (arrs1 -> a)
f Idx aenv s
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv arrs1 -> Int
forall a. acc aenv a -> Int
countA acc aenv arrs1
a
      Aforeign ArraysR a
_ asm (as -> a)
_ PreAfun acc (as -> a)
_ acc aenv as
a           -> acc aenv as -> Int
forall a. acc aenv a -> Int
countA acc aenv as
a
      Acond Exp aenv TAG
p acc aenv a
t acc aenv a
e                -> Exp aenv TAG -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv TAG
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv a -> Int
forall a. acc aenv a -> Int
countA acc aenv a
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv a -> Int
forall a. acc aenv a -> Int
countA acc aenv a
e
      -- Body and condition of the while loop may be evaluated multiple times.
      -- We multiply the usage count, as a practical solution to this. As
      -- we will check whether the count is at most 1, we will thus never
      -- inline variables used in while loops.
      Awhile PreOpenAfun acc aenv (a -> Scalar TAG)
c PreOpenAfun acc aenv (a -> a)
f acc aenv a
a               -> Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* PreOpenAfun acc aenv (a -> Scalar TAG) -> Idx aenv s -> Int
forall aenv' f. PreOpenAfun acc aenv' f -> Idx aenv' s -> Int
countAF PreOpenAfun acc aenv (a -> Scalar TAG)
c Idx aenv s
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* PreOpenAfun acc aenv (a -> a) -> Idx aenv s -> Int
forall aenv' f. PreOpenAfun acc aenv' f -> Idx aenv' s -> Int
countAF PreOpenAfun acc aenv (a -> a)
f Idx aenv s
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv a -> Int
forall a. acc aenv a -> Int
countA acc aenv a
a
      Use ArrayR (Array sh e)
_ Array sh e
_                    -> Int
0
      Unit TypeR e
_ Exp aenv e
e                   -> Exp aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv e
e
      Reshape ShapeR sh
_ Exp aenv sh
e acc aenv (Array sh' e)
a              -> Exp aenv sh -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv sh
e  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh' e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh' e)
a
      Generate ArrayR (Array sh e)
_ Exp aenv sh
e Fun aenv (sh -> e)
f             -> Exp aenv sh -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv sh
e  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Fun aenv (sh -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (sh -> e)
f
      Transform ArrayR (Array sh' b)
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
ix Fun aenv (a -> b)
f acc aenv (Array sh a)
a      -> Exp aenv sh' -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv sh'
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Fun aenv (sh' -> sh) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (sh' -> sh)
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Fun aenv (a -> b) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (a -> b)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh a) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh a)
a
      Replicate SliceIndex slix sl co sh
_ Exp aenv slix
sh acc aenv (Array sl e)
a           -> Exp aenv slix -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv slix
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sl e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sl e)
a
      Slice SliceIndex slix sl co sh
_ acc aenv (Array sh e)
a Exp aenv slix
sl               -> Exp aenv slix -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv slix
sl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e)
a
      Map TypeR e'
_ Fun aenv (e -> e')
f acc aenv (Array sh e)
a                  -> Fun aenv (e -> e') -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e')
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e)
a
      ZipWith TypeR e3
_ Fun aenv (e1 -> e2 -> e3)
f acc aenv (Array sh e1)
a1 acc aenv (Array sh e2)
a2          -> Fun aenv (e1 -> e2 -> e3) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e1 -> e2 -> e3)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e1) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e1)
a1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e2) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e2)
a2
      Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z acc aenv (Array (sh, Int) e)
a                 -> Fun aenv (e -> e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e -> e)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Maybe (Exp aenv e) -> Int
forall env e. Maybe (OpenExp env aenv e) -> Int
countME Maybe (Exp aenv e)
z Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array (sh, Int) e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array (sh, Int) e)
a
      FoldSeg IntegralType i
_ Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z acc aenv (Array (sh, Int) e)
a acc aenv (Segments i)
s          -> Fun aenv (e -> e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e -> e)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Maybe (Exp aenv e) -> Int
forall env e. Maybe (OpenExp env aenv e) -> Int
countME Maybe (Exp aenv e)
z Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array (sh, Int) e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array (sh, Int) e)
a  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Segments i) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Segments i)
s
      Scan  Direction
_ Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
z acc aenv (Array (sh, Int) e)
a              -> Fun aenv (e -> e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e -> e)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Maybe (Exp aenv e) -> Int
forall env e. Maybe (OpenExp env aenv e) -> Int
countME Maybe (Exp aenv e)
z Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array (sh, Int) e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array (sh, Int) e)
a
      Scan' Direction
_ Fun aenv (e -> e -> e)
f Exp aenv e
z acc aenv (Array (sh, Int) e)
a              -> Fun aenv (e -> e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e -> e)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Exp aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv e
z  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array (sh, Int) e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array (sh, Int) e)
a
      Permute Fun aenv (e -> e -> e)
f1 acc aenv (Array sh' e)
a1 Fun aenv (sh -> PrimMaybe sh')
f2 acc aenv (Array sh e)
a2        -> Fun aenv (e -> e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (e -> e -> e)
f1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh' e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh' e)
a1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Fun aenv (sh -> PrimMaybe sh') -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (sh -> PrimMaybe sh')
f2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e)
a2
      Backpermute ShapeR sh'
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
f acc aenv (Array sh e)
a       -> Exp aenv sh' -> Int
forall env e. OpenExp env aenv e -> Int
countE Exp aenv sh'
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Fun aenv (sh' -> sh) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (sh' -> sh)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e)
a
      Stencil StencilR sh e stencil
_ TypeR e'
_ Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
_ acc aenv (Array sh e)
a          -> Fun aenv (stencil -> e') -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (stencil -> e')
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh e) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh e)
a
      Stencil2 StencilR sh a stencil1
_ StencilR sh b stencil2
_ TypeR c
_ Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a)
_ acc aenv (Array sh a)
a1 Boundary aenv (Array sh b)
_ acc aenv (Array sh b)
a2 -> Fun aenv (stencil1 -> stencil2 -> c) -> Int
forall env f. OpenFun env aenv f -> Int
countF Fun aenv (stencil1 -> stencil2 -> c)
f  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh a) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh a)
a1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ acc aenv (Array sh b) -> Int
forall a. acc aenv a -> Int
countA acc aenv (Array sh b)
a2
      -- Collect s                 -> countS s

    countE :: OpenExp env aenv e -> Int
    countE :: OpenExp env aenv e -> Int
countE OpenExp env aenv e
exp = case OpenExp env aenv e
exp of
      Let ELeftHandSide bnd_t env env'
_ OpenExp env aenv bnd_t
bnd OpenExp env' aenv e
body             -> OpenExp env aenv bnd_t -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv bnd_t
bnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env' aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env' aenv e
body
      Evar ExpVar env e
_                     -> Int
0
      Const ScalarType e
_ e
_                  -> Int
0
      Undef ScalarType e
_                    -> Int
0
      OpenExp env aenv e
Nil                        -> Int
0
      Pair OpenExp env aenv t1
x OpenExp env aenv t2
y                   -> OpenExp env aenv t1 -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv t1
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv t2 -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv t2
y
      VecPack   VecR n s tup
_ OpenExp env aenv tup
e              -> OpenExp env aenv tup -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv tup
e
      VecUnpack VecR n s e
_ OpenExp env aenv (Vec n s)
e              -> OpenExp env aenv (Vec n s) -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv (Vec n s)
e
      IndexSlice SliceIndex slix e co sh
_ OpenExp env aenv slix
ix OpenExp env aenv sh
sh         -> OpenExp env aenv slix -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv slix
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv sh -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv sh
sh
      IndexFull SliceIndex slix sl co e
_ OpenExp env aenv slix
ix OpenExp env aenv sl
sl          -> OpenExp env aenv slix -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv slix
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv sl -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv sl
sl
      ToIndex ShapeR sh
_ OpenExp env aenv sh
sh OpenExp env aenv sh
ix            -> OpenExp env aenv sh -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv sh
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv sh -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv sh
ix
      FromIndex ShapeR e
_ OpenExp env aenv e
sh OpenExp env aenv Int
i           -> OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv e
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv Int -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv Int
i
      Case OpenExp env aenv TAG
e [(TAG, OpenExp env aenv e)]
rhs Maybe (OpenExp env aenv e)
def             -> OpenExp env aenv TAG -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv TAG
e  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv e
c | (TAG
_,OpenExp env aenv e
c) <- [(TAG, OpenExp env aenv e)]
rhs ] Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
-> (OpenExp env aenv e -> Int) -> Maybe (OpenExp env aenv e) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE Maybe (OpenExp env aenv e)
def
      Cond OpenExp env aenv TAG
p OpenExp env aenv e
t OpenExp env aenv e
e                 -> OpenExp env aenv TAG -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv TAG
p  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv e
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv e
e
      While OpenFun env aenv (e -> TAG)
p OpenFun env aenv (e -> e)
f OpenExp env aenv e
x                -> OpenFun env aenv (e -> TAG) -> Int
forall env f. OpenFun env aenv f -> Int
countF OpenFun env aenv (e -> TAG)
p  Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenFun env aenv (e -> e) -> Int
forall env f. OpenFun env aenv f -> Int
countF OpenFun env aenv (e -> e)
f Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv e
x
      PrimConst PrimConst e
_                -> Int
0
      PrimApp PrimFun (a -> e)
_ OpenExp env aenv a
x                -> OpenExp env aenv a -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv a
x
      Index ArrayVar aenv (Array dim e)
a OpenExp env aenv dim
sh                 -> ArrayVar aenv (Array dim e) -> Int
forall a. ArrayVar aenv a -> Int
countAvar ArrayVar aenv (Array dim e)
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv dim -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv dim
sh
      LinearIndex ArrayVar aenv (Array dim e)
a OpenExp env aenv Int
i            -> ArrayVar aenv (Array dim e) -> Int
forall a. ArrayVar aenv a -> Int
countAvar ArrayVar aenv (Array dim e)
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ OpenExp env aenv Int -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv Int
i
      ShapeSize ShapeR dim
_ OpenExp env aenv dim
sh             -> OpenExp env aenv dim -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv dim
sh
      Shape ArrayVar aenv (Array e e)
a
        | Bool
withShape              -> ArrayVar aenv (Array e e) -> Int
forall a. ArrayVar aenv a -> Int
countAvar ArrayVar aenv (Array e e)
a
        | Bool
otherwise              -> Int
0
      Foreign TypeR e
_ asm (x -> e)
_ Fun () (x -> e)
_ OpenExp env aenv x
e            -> OpenExp env aenv x -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv x
e
      Coerce ScalarType a
_ ScalarType e
_ OpenExp env aenv a
e               -> OpenExp env aenv a -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv a
e

    countME :: Maybe (OpenExp env aenv e) -> Int
    countME :: Maybe (OpenExp env aenv e) -> Int
countME = Int
-> (OpenExp env aenv e -> Int) -> Maybe (OpenExp env aenv e) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 OpenExp env aenv e -> Int
forall env e. OpenExp env aenv e -> Int
countE

    countA :: acc aenv a -> Int
    countA :: acc aenv a -> Int
countA = Bool -> Idx aenv s -> acc aenv a -> Int
UsesOfAcc acc
countAcc Bool
withShape Idx aenv s
idx

    countAvar :: ArrayVar aenv a -> Int
    countAvar :: ArrayVar aenv a -> Int
countAvar (Var ArrayR a
_ Idx aenv a
this) = Idx aenv a -> Int
forall a. Idx aenv a -> Int
countIdx Idx aenv a
this

    countAF :: PreOpenAfun acc aenv' f
            -> Idx aenv' s
            -> Int
    countAF :: PreOpenAfun acc aenv' f -> Idx aenv' s -> Int
countAF (Alam ALeftHandSide a aenv' aenv'
lhs PreOpenAfun acc aenv' t
f) Idx aenv' s
v = PreOpenAfun acc aenv' t -> Idx aenv' s -> Int
forall aenv' f. PreOpenAfun acc aenv' f -> Idx aenv' s -> Int
countAF PreOpenAfun acc aenv' t
f (ALeftHandSide a aenv' aenv' -> aenv' :> aenv'
forall (s :: * -> *) t env env'.
LeftHandSide s t env env' -> env :> env'
weakenWithLHS ALeftHandSide a aenv' aenv'
lhs (aenv' :> aenv') -> Idx aenv' s -> Idx aenv' s
forall env env'.
(env :> env') -> forall t'. Idx env t' -> Idx env' t'
>:> Idx aenv' s
v)
    countAF (Abody acc aenv' f
a)    Idx aenv' s
v = Bool -> Idx aenv' s -> acc aenv' f -> Int
UsesOfAcc acc
countAcc Bool
withShape Idx aenv' s
v acc aenv' f
a

    countF :: OpenFun env aenv f -> Int
    countF :: OpenFun env aenv f -> Int
countF (Lam ELeftHandSide a env env'
_ OpenFun env' aenv t
f) = OpenFun env' aenv t -> Int
forall env f. OpenFun env aenv f -> Int
countF OpenFun env' aenv t
f
    countF (Body  OpenExp env aenv f
b) = OpenExp env aenv f -> Int
forall env e. OpenExp env aenv e -> Int
countE OpenExp env aenv f
b

{--
    countS :: PreOpenSeq acc aenv senv arrs -> Int
    countS seq =
      case seq of
        Producer p s -> countP p + countS s
        Consumer c   -> countC c
        Reify _      -> 0

    countP :: Producer acc aenv senv arrs -> Int
    countP p =
      case p of
        StreamIn _           -> 0
        ToSeq _ _ a          -> countA a
        MapSeq f _           -> countAF f idx
        ChunkedMapSeq f _    -> countAF f idx
        ZipWithSeq f _ _     -> countAF f idx
        ScanSeq f e _        -> countF f + countE e

    countC :: Consumer acc aenv senv arrs -> Int
    countC c =
      case c of
        FoldSeq f e _        -> countF f + countE e
        FoldSeqFlatten f a _ -> countAF f idx + countA a
        Stuple t             -> countCT t

    countCT :: Atuple (Consumer acc aenv senv) t' -> Int
    countCT NilAtup        = 0
    countCT (SnocAtup t c) = countCT t + countC c
--}