content: {-# 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 STRING_LIT_MAX_LENGTH = 80
  
  unfloatStringLits
    :: UniqSet Name
    -> UniqFM Name ByteString
    -> [CgStgBinding]
    -> ([CgStgBinding], UniqSet Name)
  unfloatStringLits usedOnceStringLits stringLits =
    unfloatStringLits' (selectStringLitsForUnfloat usedOnceStringLits 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' stringLits allBindings = (binderWithoutChanges ++ binderWithUnfloatedStringLit, actuallyUsedStringLitNames)
    where
      (binderWithoutChanges, binderWithUnfloatedStringLitPairs) = partitionWith substituteStringLit allBindings
  
      binderWithUnfloatedStringLit = fst <$> binderWithUnfloatedStringLitPairs
      actuallyUsedStringLitNames = unionManyUniqSets (snd <$> binderWithUnfloatedStringLitPairs)
  
      substituteStringLit :: CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
      substituteStringLit x@(StgRec bnds)
        | isEmptyUniqSet names = Left x
        | otherwise = Right (StgRec bnds', names)
        where
          (bnds', names) = extractNames id $ do
            (i, rhs) <- bnds
            pure $ case processStgRhs rhs of
              Nothing -> Left (i, rhs)
              Just (rhs', names) -> Right ((i, rhs'), names)
      substituteStringLit x@(StgNonRec binder rhs)
        = maybe (Left x)
          (\(body', names) -> Right (StgNonRec binder body', names))
          (processStgRhs rhs)
  
      processStgRhs :: CgStgRhs -> Maybe (CgStgRhs, UniqSet Name)
      processStgRhs (StgRhsCon ccs dataCon mu ticks args typ)
        | isEmptyUniqSet names = Nothing
        | otherwise = Just (StgRhsCon ccs dataCon mu ticks unified typ, names)
        where
          (unified, names) = substituteArgWithNames args
      processStgRhs (StgRhsClosure fvs ccs upd bndrs body typ)
        = (\(body', names) -> (StgRhsClosure fvs ccs upd bndrs body' typ, names)) <$>
          processStgExpr body
  
      -- Recursive expressions
      processStgExpr :: CgStgExpr -> Maybe (CgStgExpr, UniqSet Name)
      processStgExpr (StgLit _) = Nothing
      processStgExpr (StgTick _ _) = Nothing
      processStgExpr (StgLet n b e) =
        case (substituteStringLit b, processStgExpr e) of
          (Left _, Nothing) -> Nothing
          (Right (b', names), Nothing) -> Just (StgLet n b' e, names)
          (Left _, Just (e', names)) -> Just (StgLet n b e', names)
          (Right (b', names), Just (e', names')) -> Just (StgLet n b' e', names `unionUniqSets` names')
      processStgExpr (StgLetNoEscape n b e) =
        case (substituteStringLit b, processStgExpr e) of
          (Left _, Nothing) -> Nothing
          (Right (b', names), Nothing) -> Just (StgLetNoEscape n b' e, names)
          (Left _, Just (e', names)) -> Just (StgLetNoEscape n b e', names)
          (Right (b', names), Just (e', names')) -> Just (StgLetNoEscape n b' e', names `unionUniqSets` names')
      -- We should keep the order: See Note [Case expression invariants]
      processStgExpr (StgCase e bndr alt_type alts) =
        case (isEmptyUniqSet names, processStgExpr e) of
          (True, Nothing) -> Nothing
          (True, Just (e', names')) -> Just (StgCase e' bndr alt_type alts, names')
          (False, Nothing) -> Just (StgCase e bndr alt_type unified, names)
          (False, Just (e', names')) -> Just (StgCase e' bndr alt_type unified, names `unionUniqSets` names')
        where
          (unified, names) = extractNames splitAlts alts
  
          splitAlts :: CgStgAlt -> Either CgStgAlt (CgStgAlt, UniqSet Name)
          splitAlts alt@(GenStgAlt con bndrs rhs) =
            case processStgExpr rhs of
              Nothing -> Left alt
              Just (alt', names) -> Right (GenStgAlt con bndrs alt', names)
  
      -- No args
      processStgExpr (StgApp _ []) = Nothing
      processStgExpr (StgConApp _ _ [] _) = Nothing
      processStgExpr (StgOpApp _ [] _) = Nothing
  
      -- Main targets. Preserving the order of args is important
      processStgExpr (StgApp fn args@(_:_))
        | isEmptyUniqSet names = Nothing
        | otherwise = Just (StgApp fn unified, names)
        where
          (unified, names) = substituteArgWithNames args
      processStgExpr (StgConApp dc n args@(_:_) tys)
        | isEmptyUniqSet names = Nothing
        | otherwise = Just (StgConApp dc n unified tys, names)
        where
          (unified, names) = substituteArgWithNames args
      processStgExpr (StgOpApp op args@(_:_) tys)
        | isEmptyUniqSet names = Nothing
        | otherwise = Just (StgOpApp op unified tys, names)
        where
          (unified, names) = substituteArgWithNames args
  
      substituteArg :: StgArg -> Either StgArg (StgArg, Name)
      substituteArg a@(StgLitArg _) = Left a
      substituteArg a@(StgVarArg i) =
        let name = idName i
        in case lookupUFM stringLits name of
          Nothing -> Left a
          Just b -> Right (StgLitArg $ LitString b, name)
  
      substituteArgWithNames = extractNames (second (second unitUniqSet) . substituteArg)
  
      extractNames :: (a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
      extractNames splitter target =
        let
          splitted = splitter <$> target
          combined = either (, emptyUniqSet) id <$> splitted
          unified = fst <$> combined
          names = unionManyUniqSets (snd <$> combined)
        in (unified, names)
  
  selectStringLitsForUnfloat :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
  selectStringLitsForUnfloat usedOnceStringLits stringLits = alwaysUnfloat `plusUFM` usedOnceUnfloat
    where
      alwaysUnfloat = alwaysUnfloatStringLits stringLits
      usedOnceUnfloat = selectUsedOnceStringLits usedOnceStringLits stringLits
  
      alwaysUnfloatStringLits :: UniqFM Name ByteString -> UniqFM Name ByteString
      alwaysUnfloatStringLits = filterUFM $ \b -> BS.length b < STRING_LIT_MAX_LENGTH
  
      selectUsedOnceStringLits :: UniqSet Name -> UniqFM Name ByteString -> UniqFM Name ByteString
      selectUsedOnceStringLits usedOnceStringLits stringLits =
        stringLits `intersectUFM` getUniqSet usedOnceStringLits
  