{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

module GHC.StgToJS.Sinker.StringsUnfloat
  ( unfloatStringLits
  )
  where

import GHC.Prelude
import GHC.Types.Unique.Set
import GHC.Types.Unique.FM
import GHC.Stg.Syntax
import GHC.Types.Id
import GHC.Types.Name
import GHC.Types.Literal
import GHC.Utils.Misc (partitionWith)

import Data.ByteString qualified as BS
import Data.ByteString (ByteString)
import Data.Bifunctor (Bifunctor (..))

-- | We suppose that every string shorter than 80 symbols is safe for sink.
-- Sinker is working on per module. It means that ALL locally defined strings
-- in a module shorter 80 symbols will be unfloated back.
pattern STRING_LIT_MAX_LENGTH :: Int
pattern $mSTRING_LIT_MAX_LENGTH :: forall {r}. Int -> ((# #) -> r) -> ((# #) -> r) -> r
$bSTRING_LIT_MAX_LENGTH :: Int
STRING_LIT_MAX_LENGTH = 80

unfloatStringLits
  :: UniqSet Name
  -> UniqFM Name ByteString
  -> [CgStgBinding]
  -> ([CgStgBinding], UniqSet Name)
unfloatStringLits :: UniqSet Name
-> UniqFM Name ByteString
-> [CgStgBinding]
-> ([CgStgBinding], UniqSet Name)
unfloatStringLits UniqSet Name
usedOnceStringLits UniqFM Name ByteString
stringLits =
  UniqFM Name ByteString
-> [CgStgBinding] -> ([CgStgBinding], UniqSet Name)
unfloatStringLits' (UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
selectStringLitsForUnfloat UniqSet Name
usedOnceStringLits UniqFM Name ByteString
stringLits)

-- | We are doing attempts to unfloat string literals back to
-- the call site. Further special JS optimizations
-- can generate more performant operations over them.
unfloatStringLits' :: UniqFM Name ByteString -> [CgStgBinding] -> ([CgStgBinding], UniqSet Name)
unfloatStringLits' :: UniqFM Name ByteString
-> [CgStgBinding] -> ([CgStgBinding], UniqSet Name)
unfloatStringLits' UniqFM Name ByteString
stringLits [CgStgBinding]
allBindings = ([CgStgBinding]
binderWithoutChanges [CgStgBinding] -> [CgStgBinding] -> [CgStgBinding]
forall a. [a] -> [a] -> [a]
++ [CgStgBinding]
binderWithUnfloatedStringLit, UniqSet Name
actuallyUsedStringLitNames)
  where
    ([CgStgBinding]
binderWithoutChanges, [(CgStgBinding, UniqSet Name)]
binderWithUnfloatedStringLitPairs) = (CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name))
-> [CgStgBinding]
-> ([CgStgBinding], [(CgStgBinding, UniqSet Name)])
forall a b c. (a -> Either b c) -> [a] -> ([b], [c])
partitionWith CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
substituteStringLit [CgStgBinding]
allBindings

    binderWithUnfloatedStringLit :: [CgStgBinding]
binderWithUnfloatedStringLit = (CgStgBinding, UniqSet Name) -> CgStgBinding
forall a b. (a, b) -> a
fst ((CgStgBinding, UniqSet Name) -> CgStgBinding)
-> [(CgStgBinding, UniqSet Name)] -> [CgStgBinding]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(CgStgBinding, UniqSet Name)]
binderWithUnfloatedStringLitPairs
    actuallyUsedStringLitNames :: UniqSet Name
actuallyUsedStringLitNames = [UniqSet Name] -> UniqSet Name
forall a. [UniqSet a] -> UniqSet a
unionManyUniqSets ((CgStgBinding, UniqSet Name) -> UniqSet Name
forall a b. (a, b) -> b
snd ((CgStgBinding, UniqSet Name) -> UniqSet Name)
-> [(CgStgBinding, UniqSet Name)] -> [UniqSet Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(CgStgBinding, UniqSet Name)]
binderWithUnfloatedStringLitPairs)

    substituteStringLit :: CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
    substituteStringLit :: CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
substituteStringLit x :: CgStgBinding
x@(StgRec [(BinderP 'CodeGen, GenStgRhs 'CodeGen)]
bnds)
      | UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names = CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
forall a b. a -> Either a b
Left CgStgBinding
x
      | Bool
otherwise = (CgStgBinding, UniqSet Name)
-> Either CgStgBinding (CgStgBinding, UniqSet Name)
forall a b. b -> Either a b
Right ([(BinderP 'CodeGen, GenStgRhs 'CodeGen)] -> CgStgBinding
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [(Id, GenStgRhs 'CodeGen)]
[(BinderP 'CodeGen, GenStgRhs 'CodeGen)]
bnds', UniqSet Name
names)
      where
        ([(Id, GenStgRhs 'CodeGen)]
bnds', UniqSet Name
names) = (Either
   (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)
 -> Either
      (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name))
-> [Either
      (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)]
-> ([(Id, GenStgRhs 'CodeGen)], UniqSet Name)
forall a x.
(a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
extractNames Either
  (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)
-> Either
     (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)
forall a. a -> a
id ([Either
    (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)]
 -> ([(Id, GenStgRhs 'CodeGen)], UniqSet Name))
-> [Either
      (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)]
-> ([(Id, GenStgRhs 'CodeGen)], UniqSet Name)
forall a b. (a -> b) -> a -> b
$ do
          (i, rhs) <- [(Id, GenStgRhs 'CodeGen)]
[(BinderP 'CodeGen, GenStgRhs 'CodeGen)]
bnds
          pure $ case processStgRhs rhs of
            Maybe (GenStgRhs 'CodeGen, UniqSet Name)
Nothing -> (Id, GenStgRhs 'CodeGen)
-> Either
     (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)
forall a b. a -> Either a b
Left (Id
i, GenStgRhs 'CodeGen
rhs)
            Just (GenStgRhs 'CodeGen
rhs', UniqSet Name
names) -> ((Id, GenStgRhs 'CodeGen), UniqSet Name)
-> Either
     (Id, GenStgRhs 'CodeGen) ((Id, GenStgRhs 'CodeGen), UniqSet Name)
forall a b. b -> Either a b
Right ((Id
i, GenStgRhs 'CodeGen
rhs'), UniqSet Name
names)
    substituteStringLit x :: CgStgBinding
x@(StgNonRec BinderP 'CodeGen
binder GenStgRhs 'CodeGen
rhs)
      = Either CgStgBinding (CgStgBinding, UniqSet Name)
-> ((GenStgRhs 'CodeGen, UniqSet Name)
    -> Either CgStgBinding (CgStgBinding, UniqSet Name))
-> Maybe (GenStgRhs 'CodeGen, UniqSet Name)
-> Either CgStgBinding (CgStgBinding, UniqSet Name)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
forall a b. a -> Either a b
Left CgStgBinding
x)
        (\(GenStgRhs 'CodeGen
body', UniqSet Name
names) -> (CgStgBinding, UniqSet Name)
-> Either CgStgBinding (CgStgBinding, UniqSet Name)
forall a b. b -> Either a b
Right (BinderP 'CodeGen -> GenStgRhs 'CodeGen -> CgStgBinding
forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec BinderP 'CodeGen
binder GenStgRhs 'CodeGen
body', UniqSet Name
names))
        (GenStgRhs 'CodeGen -> Maybe (GenStgRhs 'CodeGen, UniqSet Name)
processStgRhs GenStgRhs 'CodeGen
rhs)

    processStgRhs :: CgStgRhs -> Maybe (CgStgRhs, UniqSet Name)
    processStgRhs :: GenStgRhs 'CodeGen -> Maybe (GenStgRhs 'CodeGen, UniqSet Name)
processStgRhs (StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args Type
typ)
      | UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names = Maybe (GenStgRhs 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
      | Bool
otherwise = (GenStgRhs 'CodeGen, UniqSet Name)
-> Maybe (GenStgRhs 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> Type
-> GenStgRhs 'CodeGen
forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> Type
-> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
unified Type
typ, UniqSet Name
names)
      where
        ([StgArg]
unified, UniqSet Name
names) = [StgArg] -> ([StgArg], UniqSet Name)
substituteArgWithNames [StgArg]
args
    processStgRhs (StgRhsClosure XRhsClosure 'CodeGen
fvs CostCentreStack
ccs UpdateFlag
upd [BinderP 'CodeGen]
bndrs GenStgExpr 'CodeGen
body Type
typ)
      = (\(GenStgExpr 'CodeGen
body', UniqSet Name
names) -> (XRhsClosure 'CodeGen
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'CodeGen]
-> GenStgExpr 'CodeGen
-> Type
-> GenStgRhs 'CodeGen
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> Type
-> GenStgRhs pass
StgRhsClosure XRhsClosure 'CodeGen
fvs CostCentreStack
ccs UpdateFlag
upd [BinderP 'CodeGen]
bndrs GenStgExpr 'CodeGen
body' Type
typ, UniqSet Name
names)) ((GenStgExpr 'CodeGen, UniqSet Name)
 -> (GenStgRhs 'CodeGen, UniqSet Name))
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgRhs 'CodeGen, UniqSet Name)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr GenStgExpr 'CodeGen
body

    -- Recursive expressions
    processStgExpr :: CgStgExpr -> Maybe (CgStgExpr, UniqSet Name)
    processStgExpr :: GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr (StgLit Literal
_) = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
    processStgExpr (StgTick StgTickish
_ GenStgExpr 'CodeGen
_) = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
    processStgExpr (StgLet XLet 'CodeGen
n CgStgBinding
b GenStgExpr 'CodeGen
e) =
      case (CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
substituteStringLit CgStgBinding
b, GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr GenStgExpr 'CodeGen
e) of
        (Left CgStgBinding
_, Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
        (Right (CgStgBinding
b', UniqSet Name
names), Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLet 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet 'CodeGen
n CgStgBinding
b' GenStgExpr 'CodeGen
e, UniqSet Name
names)
        (Left CgStgBinding
_, Just (GenStgExpr 'CodeGen
e', UniqSet Name
names)) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLet 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet 'CodeGen
n CgStgBinding
b GenStgExpr 'CodeGen
e', UniqSet Name
names)
        (Right (CgStgBinding
b', UniqSet Name
names), Just (GenStgExpr 'CodeGen
e', UniqSet Name
names')) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLet 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet 'CodeGen
n CgStgBinding
b' GenStgExpr 'CodeGen
e', UniqSet Name
names UniqSet Name -> UniqSet Name -> UniqSet Name
forall a. UniqSet a -> UniqSet a -> UniqSet a
`unionUniqSets` UniqSet Name
names')
    processStgExpr (StgLetNoEscape XLetNoEscape 'CodeGen
n CgStgBinding
b GenStgExpr 'CodeGen
e) =
      case (CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
substituteStringLit CgStgBinding
b, GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr GenStgExpr 'CodeGen
e) of
        (Left CgStgBinding
_, Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
        (Right (CgStgBinding
b', UniqSet Name
names), Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLetNoEscape 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape 'CodeGen
n CgStgBinding
b' GenStgExpr 'CodeGen
e, UniqSet Name
names)
        (Left CgStgBinding
_, Just (GenStgExpr 'CodeGen
e', UniqSet Name
names)) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLetNoEscape 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape 'CodeGen
n CgStgBinding
b GenStgExpr 'CodeGen
e', UniqSet Name
names)
        (Right (CgStgBinding
b', UniqSet Name
names), Just (GenStgExpr 'CodeGen
e', UniqSet Name
names')) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (XLetNoEscape 'CodeGen
-> CgStgBinding -> GenStgExpr 'CodeGen -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape 'CodeGen
n CgStgBinding
b' GenStgExpr 'CodeGen
e', UniqSet Name
names UniqSet Name -> UniqSet Name -> UniqSet Name
forall a. UniqSet a -> UniqSet a -> UniqSet a
`unionUniqSets` UniqSet Name
names')
    -- We should keep the order: See Note [Case expression invariants]
    processStgExpr (StgCase GenStgExpr 'CodeGen
e BinderP 'CodeGen
bndr AltType
alt_type [GenStgAlt 'CodeGen]
alts) =
      case (UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names, GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr GenStgExpr 'CodeGen
e) of
        (Bool
True, Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
        (Bool
True, Just (GenStgExpr 'CodeGen
e', UniqSet Name
names')) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (GenStgExpr 'CodeGen
-> BinderP 'CodeGen
-> AltType
-> [GenStgAlt 'CodeGen]
-> GenStgExpr 'CodeGen
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'CodeGen
e' BinderP 'CodeGen
bndr AltType
alt_type [GenStgAlt 'CodeGen]
alts, UniqSet Name
names')
        (Bool
False, Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (GenStgExpr 'CodeGen
-> BinderP 'CodeGen
-> AltType
-> [GenStgAlt 'CodeGen]
-> GenStgExpr 'CodeGen
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'CodeGen
e BinderP 'CodeGen
bndr AltType
alt_type [GenStgAlt 'CodeGen]
unified, UniqSet Name
names)
        (Bool
False, Just (GenStgExpr 'CodeGen
e', UniqSet Name
names')) -> (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (GenStgExpr 'CodeGen
-> BinderP 'CodeGen
-> AltType
-> [GenStgAlt 'CodeGen]
-> GenStgExpr 'CodeGen
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'CodeGen
e' BinderP 'CodeGen
bndr AltType
alt_type [GenStgAlt 'CodeGen]
unified, UniqSet Name
names UniqSet Name -> UniqSet Name -> UniqSet Name
forall a. UniqSet a -> UniqSet a -> UniqSet a
`unionUniqSets` UniqSet Name
names')
      where
        ([GenStgAlt 'CodeGen]
unified, UniqSet Name
names) = (GenStgAlt 'CodeGen
 -> Either (GenStgAlt 'CodeGen) (GenStgAlt 'CodeGen, UniqSet Name))
-> [GenStgAlt 'CodeGen] -> ([GenStgAlt 'CodeGen], UniqSet Name)
forall a x.
(a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
extractNames GenStgAlt 'CodeGen
-> Either (GenStgAlt 'CodeGen) (GenStgAlt 'CodeGen, UniqSet Name)
splitAlts [GenStgAlt 'CodeGen]
alts

        splitAlts :: CgStgAlt -> Either CgStgAlt (CgStgAlt, UniqSet Name)
        splitAlts :: GenStgAlt 'CodeGen
-> Either (GenStgAlt 'CodeGen) (GenStgAlt 'CodeGen, UniqSet Name)
splitAlts alt :: GenStgAlt 'CodeGen
alt@(GenStgAlt AltCon
con [BinderP 'CodeGen]
bndrs GenStgExpr 'CodeGen
rhs) =
          case GenStgExpr 'CodeGen -> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
processStgExpr GenStgExpr 'CodeGen
rhs of
            Maybe (GenStgExpr 'CodeGen, UniqSet Name)
Nothing -> GenStgAlt 'CodeGen
-> Either (GenStgAlt 'CodeGen) (GenStgAlt 'CodeGen, UniqSet Name)
forall a b. a -> Either a b
Left GenStgAlt 'CodeGen
alt
            Just (GenStgExpr 'CodeGen
alt', UniqSet Name
names) -> (GenStgAlt 'CodeGen, UniqSet Name)
-> Either (GenStgAlt 'CodeGen) (GenStgAlt 'CodeGen, UniqSet Name)
forall a b. b -> Either a b
Right (AltCon
-> [BinderP 'CodeGen] -> GenStgExpr 'CodeGen -> GenStgAlt 'CodeGen
forall (pass :: StgPass).
AltCon -> [BinderP pass] -> GenStgExpr pass -> GenStgAlt pass
GenStgAlt AltCon
con [BinderP 'CodeGen]
bndrs GenStgExpr 'CodeGen
alt', UniqSet Name
names)

    -- No args
    processStgExpr (StgApp Id
_ []) = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
    processStgExpr (StgConApp DataCon
_ ConstructorNumber
_ [] [[PrimRep]]
_) = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
    processStgExpr (StgOpApp StgOp
_ [] Type
_) = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing

    -- Main targets. Preserving the order of args is important
    processStgExpr (StgApp Id
fn args :: [StgArg]
args@(StgArg
_:[StgArg]
_))
      | UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
      | Bool
otherwise = (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (Id -> [StgArg] -> GenStgExpr 'CodeGen
forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
fn [StgArg]
unified, UniqSet Name
names)
      where
        ([StgArg]
unified, UniqSet Name
names) = [StgArg] -> ([StgArg], UniqSet Name)
substituteArgWithNames [StgArg]
args
    processStgExpr (StgConApp DataCon
dc ConstructorNumber
n args :: [StgArg]
args@(StgArg
_:[StgArg]
_) [[PrimRep]]
tys)
      | UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
      | Bool
otherwise = (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (DataCon
-> ConstructorNumber
-> [StgArg]
-> [[PrimRep]]
-> GenStgExpr 'CodeGen
forall (pass :: StgPass).
DataCon
-> ConstructorNumber -> [StgArg] -> [[PrimRep]] -> GenStgExpr pass
StgConApp DataCon
dc ConstructorNumber
n [StgArg]
unified [[PrimRep]]
tys, UniqSet Name
names)
      where
        ([StgArg]
unified, UniqSet Name
names) = [StgArg] -> ([StgArg], UniqSet Name)
substituteArgWithNames [StgArg]
args
    processStgExpr (StgOpApp StgOp
op args :: [StgArg]
args@(StgArg
_:[StgArg]
_) Type
tys)
      | UniqSet Name -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet Name
names = Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. Maybe a
Nothing
      | Bool
otherwise = (GenStgExpr 'CodeGen, UniqSet Name)
-> Maybe (GenStgExpr 'CodeGen, UniqSet Name)
forall a. a -> Maybe a
Just (StgOp -> [StgArg] -> Type -> GenStgExpr 'CodeGen
forall (pass :: StgPass).
StgOp -> [StgArg] -> Type -> GenStgExpr pass
StgOpApp StgOp
op [StgArg]
unified Type
tys, UniqSet Name
names)
      where
        ([StgArg]
unified, UniqSet Name
names) = [StgArg] -> ([StgArg], UniqSet Name)
substituteArgWithNames [StgArg]
args

    substituteArg :: StgArg -> Either StgArg (StgArg, Name)
    substituteArg :: StgArg -> Either StgArg (StgArg, Name)
substituteArg a :: StgArg
a@(StgLitArg Literal
_) = StgArg -> Either StgArg (StgArg, Name)
forall a b. a -> Either a b
Left StgArg
a
    substituteArg a :: StgArg
a@(StgVarArg Id
i) =
      let name :: Name
name = Id -> Name
idName Id
i
      in case UniqFM Name ByteString -> Name -> Maybe ByteString
forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
lookupUFM UniqFM Name ByteString
stringLits Name
name of
        Maybe ByteString
Nothing -> StgArg -> Either StgArg (StgArg, Name)
forall a b. a -> Either a b
Left StgArg
a
        Just ByteString
b -> (StgArg, Name) -> Either StgArg (StgArg, Name)
forall a b. b -> Either a b
Right (Literal -> StgArg
StgLitArg (Literal -> StgArg) -> Literal -> StgArg
forall a b. (a -> b) -> a -> b
$ ByteString -> Literal
LitString ByteString
b, Name
name)

    substituteArgWithNames :: [StgArg] -> ([StgArg], UniqSet Name)
substituteArgWithNames = (StgArg -> Either StgArg (StgArg, UniqSet Name))
-> [StgArg] -> ([StgArg], UniqSet Name)
forall a x.
(a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
extractNames (((StgArg, Name) -> (StgArg, UniqSet Name))
-> Either StgArg (StgArg, Name)
-> Either StgArg (StgArg, UniqSet Name)
forall b c a. (b -> c) -> Either a b -> Either a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Name -> UniqSet Name) -> (StgArg, Name) -> (StgArg, UniqSet Name)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Name -> UniqSet Name
forall a. Uniquable a => a -> UniqSet a
unitUniqSet) (Either StgArg (StgArg, Name)
 -> Either StgArg (StgArg, UniqSet Name))
-> (StgArg -> Either StgArg (StgArg, Name))
-> StgArg
-> Either StgArg (StgArg, UniqSet Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StgArg -> Either StgArg (StgArg, Name)
substituteArg)

    extractNames :: (a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
    extractNames :: forall a x.
(a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
extractNames a -> Either x (x, UniqSet Name)
splitter [a]
target =
      let
        splitted :: [Either x (x, UniqSet Name)]
splitted = a -> Either x (x, UniqSet Name)
splitter (a -> Either x (x, UniqSet Name))
-> [a] -> [Either x (x, UniqSet Name)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a]
target
        combined :: [(x, UniqSet Name)]
combined = (x -> (x, UniqSet Name))
-> ((x, UniqSet Name) -> (x, UniqSet Name))
-> Either x (x, UniqSet Name)
-> (x, UniqSet Name)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (, UniqSet Name
forall a. UniqSet a
emptyUniqSet) (x, UniqSet Name) -> (x, UniqSet Name)
forall a. a -> a
id (Either x (x, UniqSet Name) -> (x, UniqSet Name))
-> [Either x (x, UniqSet Name)] -> [(x, UniqSet Name)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Either x (x, UniqSet Name)]
splitted
        unified :: [x]
unified = (x, UniqSet Name) -> x
forall a b. (a, b) -> a
fst ((x, UniqSet Name) -> x) -> [(x, UniqSet Name)] -> [x]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(x, UniqSet Name)]
combined
        names :: UniqSet Name
names = [UniqSet Name] -> UniqSet Name
forall a. [UniqSet a] -> UniqSet a
unionManyUniqSets ((x, UniqSet Name) -> UniqSet Name
forall a b. (a, b) -> b
snd ((x, UniqSet Name) -> UniqSet Name)
-> [(x, UniqSet Name)] -> [UniqSet Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(x, UniqSet Name)]
combined)
      in ([x]
unified, UniqSet Name
names)

selectStringLitsForUnfloat :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
selectStringLitsForUnfloat :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
selectStringLitsForUnfloat UniqSet Name
usedOnceStringLits UniqFM Name ByteString
stringLits = UniqFM Name ByteString
alwaysUnfloat UniqFM Name ByteString
-> UniqFM Name ByteString -> UniqFM Name ByteString
forall {k} (key :: k) elt.
UniqFM key elt -> UniqFM key elt -> UniqFM key elt
`plusUFM` UniqFM Name ByteString
usedOnceUnfloat
  where
    alwaysUnfloat :: UniqFM Name ByteString
alwaysUnfloat = UniqFM Name ByteString -> UniqFM Name ByteString
alwaysUnfloatStringLits UniqFM Name ByteString
stringLits
    usedOnceUnfloat :: UniqFM Name ByteString
usedOnceUnfloat = UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
selectUsedOnceStringLits UniqSet Name
usedOnceStringLits UniqFM Name ByteString
stringLits

    alwaysUnfloatStringLits :: UniqFM Name ByteString -> UniqFM Name ByteString
    alwaysUnfloatStringLits :: UniqFM Name ByteString -> UniqFM Name ByteString
alwaysUnfloatStringLits = (ByteString -> Bool)
-> UniqFM Name ByteString -> UniqFM Name ByteString
forall {k} elt (key :: k).
(elt -> Bool) -> UniqFM key elt -> UniqFM key elt
filterUFM ((ByteString -> Bool)
 -> UniqFM Name ByteString -> UniqFM Name ByteString)
-> (ByteString -> Bool)
-> UniqFM Name ByteString
-> UniqFM Name ByteString
forall a b. (a -> b) -> a -> b
$ \ByteString
b -> ByteString -> Int
BS.length ByteString
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
STRING_LIT_MAX_LENGTH

    selectUsedOnceStringLits :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
    selectUsedOnceStringLits :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
selectUsedOnceStringLits UniqSet Name
usedOnceStringLits UniqFM Name ByteString
stringLits =
      UniqFM Name ByteString
stringLits UniqFM Name ByteString
-> UniqFM Name Name -> UniqFM Name ByteString
forall {k} (key :: k) elt1 elt2.
UniqFM key elt1 -> UniqFM key elt2 -> UniqFM key elt1
`intersectUFM` UniqSet Name -> UniqFM Name Name
forall a. UniqSet a -> UniqFM a a
getUniqSet UniqSet Name
usedOnceStringLits