<commit_before>{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- |
--
-- Perform simple hoisting of let-bindings, moving them out as far as
-- possible (particularly outside loops).
--
module L0C.EnablingOpts.Hoisting
  ( transformProg )
  where

import Control.Applicative
import Control.Arrow (first)
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.RWS

import Data.Graph
import Data.List
import Data.Loc
import Data.Maybe
import qualified Data.Map as M
import Data.Ord
import qualified Data.Set as S

import L0C.L0
import L0C.Substitute
import L0C.FreshNames

data BindNeed = LoopBind TupIdent Exp Ident Exp Exp
              | LetBind TupIdent Exp [(TupIdent, Exp)]
              | LetWithBind Certificates Ident Ident [Exp] Exp
                deriving (Show, Eq, Ord)

type NeedSet = S.Set BindNeed

asTail :: BindNeed -> Exp
asTail (LoopBind mergepat mergeexp i bound loopbody) =
  DoLoop mergepat mergeexp i bound loopbody (TupLit [] loc) loc
    where loc = srclocOf mergepat
asTail (LetBind pat e _) =
  LetPat pat e (TupLit [] loc) loc
    where loc = srclocOf pat
asTail (LetWithBind cs dest src is ve) =
  LetWith cs dest src is ve (Var dest) $ srclocOf dest

requires :: BindNeed -> S.Set VName
requires (LetBind pat e ((altpat,alte):alts)) =
  requires (LetBind pat e []) <> requires (LetBind altpat alte alts)
requires bnd = S.map identName $ freeInExp $ asTail bnd

provides :: BindNeed -> S.Set VName
provides (LoopBind mergepat _ _ _ _) = patNames mergepat
provides (LetBind pat _ _) = patNames pat
provides (LetWithBind _ dest _ _ _) = S.singleton $ identName dest

data Need = Need { needBindings :: NeedSet
                 }

instance Monoid Need where
  Need b1 `mappend` Need b2 = Need $ b1 <> b2
  mempty = Need S.empty

data ShapeBinding = DimSizes [Maybe Exp]
                    deriving (Show, Eq)

type ShapeMap = M.Map Ident [ShapeBinding]

data Env = Env { envBindings :: ShapeMap
               }

emptyEnv :: Env
emptyEnv = Env {
             envBindings = M.empty
           }

cartesian :: [[a]] -> [[a]]
cartesian [] = []
cartesian [x] = map (:[]) x
cartesian (x:xs) = [ x' : xs' | xs' <- cartesian xs, x' <- x ]

varExp :: Exp -> Maybe Ident
varExp (Var k) = Just k
varExp _       = Nothing

vars :: [Exp] -> [Ident]
vars = mapMaybe varExp

isArrayIdent :: Ident -> Bool
isArrayIdent idd = case identType idd of
                     Array {} -> True
                     _        -> False

lookupShapeBindings :: Ident -> ShapeMap -> [ShapeBinding]
lookupShapeBindings = delve S.empty
  where
    delve s k m | k `S.member` s = []
                | otherwise =
                  case M.lookup k m of
                    Nothing -> []
                    Just shs -> shs ++ concatMap (recurse s m) shs
    recurse s m (DimSizes sz) =
      map DimSizes $ filter (not . all (==Nothing)) $
          cartesian $ map inspect sz
      where inspect (Just (Size _ i (Var k') _)) =
              case delve (k' `S.insert` s) k' m of
                [] -> [Nothing]
                l  -> map (fill i) l
            inspect _ = [Nothing]
    fill i (DimSizes sz) = case drop i sz of
                             e:_ -> e
                             []  -> Nothing

newtype HoistM a = HoistM (RWS<commit_msg>Fix recursive size tracking in hoister.<commit_after>{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- |
--
-- Perform simple hoisting of let-bindings, moving them out as far as
-- possible (particularly outside loops).
--
module L0C.EnablingOpts.Hoisting
  ( transformProg )
  where

import Control.Applicative
import Control.Arrow (first)
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.RWS

import Data.Graph
import Data.List
import Data.Loc
import Data.Maybe
import qualified Data.Map as M
import Data.Ord
import qualified Data.Set as S

import L0C.L0
import L0C.Substitute
import L0C.FreshNames

data BindNeed = LoopBind TupIdent Exp Ident Exp Exp
              | LetBind TupIdent Exp [(TupIdent, Exp)]
              | LetWithBind Certificates Ident Ident [Exp] Exp
                deriving (Show, Eq, Ord)

type NeedSet = S.Set BindNeed

asTail :: BindNeed -> Exp
asTail (LoopBind mergepat mergeexp i bound loopbody) =
  DoLoop mergepat mergeexp i bound loopbody (TupLit [] loc) loc
    where loc = srclocOf mergepat
asTail (LetBind pat e _) =
  LetPat pat e (TupLit [] loc) loc
    where loc = srclocOf pat
asTail (LetWithBind cs dest src is ve) =
  LetWith cs dest src is ve (Var dest) $ srclocOf dest

requires :: BindNeed -> S.Set VName
requires (LetBind pat e ((altpat,alte):alts)) =
  requires (LetBind pat e []) <> requires (LetBind altpat alte alts)
requires bnd = S.map identName $ freeInExp $ asTail bnd

provides :: BindNeed -> S.Set VName
provides (LoopBind mergepat _ _ _ _) = patNames mergepat
provides (LetBind pat _ _) = patNames pat
provides (LetWithBind _ dest _ _ _) = S.singleton $ identName dest

data Need = Need { needBindings :: NeedSet
                 }

instance Monoid Need where
  Need b1 `mappend` Need b2 = Need $ b1 <> b2
  mempty = Need S.empty

data ShapeBinding = DimSizes [Maybe Exp]
                    deriving (Show, Eq)

type ShapeMap = M.Map Ident [ShapeBinding]

data Env = Env { envBindings :: ShapeMap
               }

emptyEnv :: Env
emptyEnv = Env {
             envBindings = M.empty
           }

cartesian :: [[a]] -> [[a]]
cartesian [] = []
cartesian [x] = map (:[]) x
cartesian (x:xs) = [ x' : xs' | xs' <- cartesian xs, x' <- x ]

varExp :: Exp -> Maybe Ident
varExp (Var k) = Just k
varExp _       = Nothing

vars :: [Exp] -> [Ident]
vars = mapMaybe varExp

isArrayIdent :: Ident -> Bool
isArrayIdent idd = case identType idd of
                     Array {} -> True
                     _        -> False

lookupShapeBindings :: Ident -> ShapeMap -> [ShapeBinding]
lookupShapeBindings idd m = delve S.empty idd
  where
    delve s k | k `S.member` s = []
              | otherwise =
                case M.lookup k m of
                  Nothing -> []
                  Just shs -> shs ++ concatMap (recurse $ k `S.insert` s) shs
    recurse s (DimSizes sz) =
      map DimSizes $ filter (not . all (==Nothing)) $
          cartesian $ map inspect sz
      where inspect (Just (Size _ i (Var k') _)) =
              case delve s k' of
                [] -> [Nothing]
                l  -> map (fill i) l
            inspect _ = [Nothing]
    fill i (DimSizes sz) = case drop i sz of
                             e:_ -> e
                             []  -> Nothing

newtype HoistM a = HoistM (RWS
{'fraction_non_alphanumeric': 0.09771246780790789, 'fraction_numerical': 0.0027268595667323133, 'mean_word_length': 2.292768079800499, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4311071', 'n_tokens_mistral': 2128, 'n_tokens_neox': 2075, 'n_words': 727}