- ar' = go1 other `max` 0
-
- go1 (Var v) = idArity v
- go1 (Lam x e) | isId x = go1 e + 1
- | otherwise = go1 e
- go1 (Note n e) | ok_note n = go1 e
- go1 (App f (Type _)) = go1 f
- go1 (App f a) | exprIsCheap a = go1 f - 1
- go1 (Case scrut _ alts)
- | exprIsCheap scrut = min_zero [go1 rhs | (_,_,rhs) <- alts]
- go1 (Let b e)
- | all exprIsCheap (rhssOfBind b) = go1 e
-
- go1 other = 0
+ ar' = length (go1 other)
+
+ go1 :: CoreExpr -> [Bool]
+ -- (go1 e) = [b1,..,bn]
+ -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
+ -- where bi is True <=> the lambda is one-shot
+
+ go1 (Note n e) | ok_note n = go1 e
+ go1 (Var v) = replicate (idArity v) False -- When the type of the Id
+ -- encodes one-shot-ness, use
+ -- th iinfo here
+
+ -- Lambdas; increase arity
+ go1 (Lam x e) | isId x = isOneShotLambda x : go1 e
+ | otherwise = go1 e
+
+ -- Applications; decrease arity
+ go1 (App f (Type _)) = go1 f
+ go1 (App f a) = case go1 f of
+ (one_shot : xs) | one_shot || exprIsCheap a -> xs
+ other -> []
+
+ -- Case/Let; keep arity if either the expression is cheap
+ -- or it's a 1-shot lambda
+ go1 (Case scrut _ alts) = case foldr1 (zipWith (&&)) [go1 rhs | (_,_,rhs) <- alts] of
+ xs@(one_shot : _) | one_shot || exprIsCheap scrut -> xs
+ other -> []
+ go1 (Let b e) = case go1 e of
+ xs@(one_shot : _) | one_shot || all exprIsCheap (rhssOfBind b) -> xs
+ other -> []
+
+ go1 other = []