<commit_before>{-# LANGUAGE ViewPatterns, TupleSections, PatternGuards, BangPatterns, RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module Supercompile.Drive (supercompile) where

import Supercompile.Match
import Supercompile.Split

import Core.FreeVars
import Core.Renaming
import Core.Syntax
import Core.Tag

import Evaluator.Evaluate
import Evaluator.FreeVars
import Evaluator.Residualise
import Evaluator.Syntax

import Size.Deeds

import Termination.TagBag
import Termination.TagGraph
import Termination.TagSet
import Termination.Terminate
import Termination.Generaliser

import Name
import Renaming
import StaticFlags
import Utilities

import Algebra.Lattice

import qualified Data.Map as M
import Data.Ord
import qualified Data.Set as S
import qualified Data.IntSet as IS
import Data.Tree


wQO :: WQO State Generaliser
wQO | not tERMINATION_CHECK                        = postcomp (const generaliseNothing) unsafeNever
    | otherwise = case tAG_COLLECTION of TagBag   -> embedWithTagBags
                                         TagGraph -> embedWithTagGraphs
                                         TagSet   -> embedWithTagSets

supercompile :: Term -> Term
supercompile e = traceRender ("all input FVs", input_fvs) $ fVedTermToTerm $ runScpM $ fmap snd $ sc (mkHistory (extra wQO)) (deeds, state)
  where input_fvs = annedTermFreeVars anned_e
        state = (Heap (setToMap Environmental input_fvs) reduceIdSupply, [], (mkIdentityRenaming $ S.toList input_fvs, anned_e))
        anned_e = toAnnedTerm e
        
        deeds = mkDeeds (bLOAT_FACTOR - 1) (t, pPrint . rb)
        
        (t, rb) = extractDeeds (\f e -> let (ts, rb) = f (annee e)
                                        in (Node (annedTag e) ts, \(Node unc ts') -> Counted unc (rb ts'))) anned_e
        
        extractDeeds :: (forall a b.    (a        -> ([Tree Tag], [Tree Int] -> b))
                                     -> Anned a   -> (Tree Tag,   Tree Int   -> Counted b))
                     -> AnnedTerm -> (Tree Tag, Tree Int -> CountedTerm)
        extractDeeds rec = term
          where 
            var = rec var'
            var' x = ([], \[] -> x)
            
            term = rec term'
            term' e = case e of
              Var x              -> ([], \[] -> Var x)
              Value (Lambda x e) -> ([t], \[t'] -> Value (Lambda x (rb t')))
                where (t, rb) = term e
              Value (Data dc xs) -> ([], \[] -> Value (Data dc xs))
              Value (Literal l)  -> ([], \[] -> Value (Literal l))
              App e x            -> ([t1, t2], \[t1', t2'] -> App (rb1 t1') (rb2 t2'))
                where (t1, rb1) = term e
                      (t2, rb2) = var x
              PrimOp pop es      -> (ts, \ts' -> PrimOp pop (zipWith ($) rbs ts'))
                where (ts, rbs) = unzip (map term es)
              Case e (unzip -> (alt_cons, alt_es)) -> (t : ts, \(t':ts') -> Case (rb t') (alt_cons `zip` zipWith ($) rbs ts'))
                where (t, rb)   = term e
                      (ts, rbs) = unzip (map term alt_es)
              LetRec (unzip -> (xs, es)) e         -> (t : ts, \(t':ts') -> LetRec (xs `zip` zipWith ($) rbs ts') (rb t'))
                where (t, rb)   = term e
                      (ts, rbs) = unzip (map term es)


--
-- == Bounded multi-step reduction ==
--

gc :: (Deeds, State) -> (Deeds, State)
gc (deeds, (Heap h ids, k, in_e)) = (M.fold (flip releaseHeapBindingDeeds) deeds h_dead,
                                     (Heap h_reachable ids, k, in_e))
  where
    (h_dead, h_reachable) = M.mapEitherWithKey (\x' hb -> maybe (Left hb) (\why_live -> Right $ demoteHeapBinding why_live hb) (x' `whyLive` reachable)) h
    reachable = lfpFrom (mkConcreteLiveness $ snd $ stackFreeVars k $ inFreeVars annedTermFreeVars in_e) go
    go live = M.foldWithKey (\x' hb live -> maybe live (\why_live -> live `plusLiveness` heapBindingLiveness (demoteHeapBinding why_live hb)) (x' `whyLive` live)) live h
    
    -- TODO: fix this comment. It used to be attached to the update frame logic in the evaluator.
    --
    -- If we can GC the update frame (because it can't be referred to in the continuation) then:
    --  1) We don't have to actually update the heap or even claim a new deed
    --  2) We make the supercompiler less likely to terminate, because doing so tends to reduce TagBag sizes
    --
    -- NB: to prevent incorrectly garbage collecting bindings from the enclosing heap when we have speculation on,
    -- we pass around an extra "live set" of parts of the heap that might be referred to later on
    --
    -- TODO: make finding FVs much cheaper (i.e. memoise it in the syntax functor construction)
    -- TODO: could GC cycles as well (i.e. don't consider stuff from the Heap that was only referred to by the thing being removed as "GC roots")

speculate :: ((Deeds, State) -> (Deeds, State))
          -> (Deeds, State) -> (Deeds, State)
speculate reduce = snd . go (0 :: Int) (mkHistory wQO) emptyLosers
  where
    go depth hist losers (deeds, state) = case terminate hist state of
        Continue hist' | sPECULATION -> continue depth hist' losers (deeds, state)
        _                            -> (losers, (deeds, state))
    
    continue depth hist losers (deeds, state@(Heap h _, _, _)) = (losers', (deeds'', (Heap (h'_winners' `M.union` h'_losers) ids'', k, in_e)))
      where
        (deeds', (Heap h' ids', k, in_e)) = reduce (deeds, state)
        
        -- It is *very important* that we prevent speculation from forcing heap-carried things that have proved
        -- to be losers in the past. This prevents us discovering speculation failure of some head binding, and
        -- then forcing several thunks that are each strict in it. This change was motivated by DigitsOfE2 -- it
        -- speeds up supercompilation of that benchmark massively.
        (h'_losers, h'_winners) | sPECULATE_ON_LOSERS = (M.empty, h')
                                | otherwise           = M.partition (\hb -> maybe False (`IS.member` losers) (heapBindingTag hb)) h'
        
        -- NB: It is important that we accumulate losers across "go" invocations in a state-monady kind of way, or DigitsOfE2 blows out
        -- even more than normal (it still takes 22s with this change).
        -- TODO: I suspect we should accumulate Losers across the boundary of speculate as well
        (deeds'', Heap h'_winners' ids'', losers') = M.foldWithKey speculate_one (deeds', Heap h'_winners ids', losers) (h'_winners M.\\ h)
        speculate_one x' (Concrete in_e) (deeds, Heap h'_winners ids, losers)<commit_msg>Comments about garbage collection only (TODOs)<commit_after>{-# LANGUAGE ViewPatterns, TupleSections, PatternGuards, BangPatterns, RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module Supercompile.Drive (supercompile) where

import Supercompile.Match
import Supercompile.Split

import Core.FreeVars
import Core.Renaming
import Core.Syntax
import Core.Tag

import Evaluator.Evaluate
import Evaluator.FreeVars
import Evaluator.Residualise
import Evaluator.Syntax

import Size.Deeds

import Termination.TagBag
import Termination.TagGraph
import Termination.TagSet
import Termination.Terminate
import Termination.Generaliser

import Name
import Renaming
import StaticFlags
import Utilities

import Algebra.Lattice

import qualified Data.Map as M
import Data.Ord
import qualified Data.Set as S
import qualified Data.IntSet as IS
import Data.Tree


wQO :: WQO State Generaliser
wQO | not tERMINATION_CHECK                        = postcomp (const generaliseNothing) unsafeNever
    | otherwise = case tAG_COLLECTION of TagBag   -> embedWithTagBags
                                         TagGraph -> embedWithTagGraphs
                                         TagSet   -> embedWithTagSets

supercompile :: Term -> Term
supercompile e = traceRender ("all input FVs", input_fvs) $ fVedTermToTerm $ runScpM $ fmap snd $ sc (mkHistory (extra wQO)) (deeds, state)
  where input_fvs = annedTermFreeVars anned_e
        state = (Heap (setToMap Environmental input_fvs) reduceIdSupply, [], (mkIdentityRenaming $ S.toList input_fvs, anned_e))
        anned_e = toAnnedTerm e
        
        deeds = mkDeeds (bLOAT_FACTOR - 1) (t, pPrint . rb)
        
        (t, rb) = extractDeeds (\f e -> let (ts, rb) = f (annee e)
                                        in (Node (annedTag e) ts, \(Node unc ts') -> Counted unc (rb ts'))) anned_e
        
        extractDeeds :: (forall a b.    (a        -> ([Tree Tag], [Tree Int] -> b))
                                     -> Anned a   -> (Tree Tag,   Tree Int   -> Counted b))
                     -> AnnedTerm -> (Tree Tag, Tree Int -> CountedTerm)
        extractDeeds rec = term
          where 
            var = rec var'
            var' x = ([], \[] -> x)
            
            term = rec term'
            term' e = case e of
              Var x              -> ([], \[] -> Var x)
              Value (Lambda x e) -> ([t], \[t'] -> Value (Lambda x (rb t')))
                where (t, rb) = term e
              Value (Data dc xs) -> ([], \[] -> Value (Data dc xs))
              Value (Literal l)  -> ([], \[] -> Value (Literal l))
              App e x            -> ([t1, t2], \[t1', t2'] -> App (rb1 t1') (rb2 t2'))
                where (t1, rb1) = term e
                      (t2, rb2) = var x
              PrimOp pop es      -> (ts, \ts' -> PrimOp pop (zipWith ($) rbs ts'))
                where (ts, rbs) = unzip (map term es)
              Case e (unzip -> (alt_cons, alt_es)) -> (t : ts, \(t':ts') -> Case (rb t') (alt_cons `zip` zipWith ($) rbs ts'))
                where (t, rb)   = term e
                      (ts, rbs) = unzip (map term alt_es)
              LetRec (unzip -> (xs, es)) e         -> (t : ts, \(t':ts') -> LetRec (xs `zip` zipWith ($) rbs ts') (rb t'))
                where (t, rb)   = term e
                      (ts, rbs) = unzip (map term es)


--
-- == Bounded multi-step reduction ==
--

-- TODO: have the garbage collector collapse indirections to indirections (but unlike GHC, not further!)
gc :: (Deeds, State) -> (Deeds, State)
gc (deeds, (Heap h ids, k, in_e)) = (M.fold (flip releaseHeapBindingDeeds) deeds h_dead,
                                     (Heap h_reachable ids, k, in_e))
  where
    (h_dead, h_reachable) = M.mapEitherWithKey (\x' hb -> maybe (Left hb) (\why_live -> Right $ demoteHeapBinding why_live hb) (x' `whyLive` reachable)) h
    reachable = lfpFrom (mkConcreteLiveness $ snd $ stackFreeVars k $ inFreeVars annedTermFreeVars in_e) go
    go live = M.foldWithKey (\x' hb live -> maybe live (\why_live -> live `plusLiveness` heapBindingLiveness (demoteHeapBinding why_live hb)) (x' `whyLive` live)) live h
    
    -- TODO: fix this comment. It used to be attached to the update frame logic in the evaluator.
    --
    -- If we can GC the update frame (because it can't be referred to in the continuation) then:
    --  1) We don't have to actually update the heap or even claim a new deed
    --  2) We make the supercompiler less likely to terminate, because doing so tends to reduce TagBag sizes
    --
    -- NB: to prevent incorrectly garbage collecting bindings from the enclosing heap when we have speculation on,
    -- we pass around an extra "live set" of parts of the heap that might be referred to later on
    --
    -- TODO: make finding FVs much cheaper (i.e. memoise it in the syntax functor construction)
    -- TODO: could GC cycles as well (i.e. don't consider stuff from the Heap that was only referred to by the thing being removed as "GC roots")

speculate :: ((Deeds, State) -> (Deeds, State))
          -> (Deeds, State) -> (Deeds, State)
speculate reduce = snd . go (0 :: Int) (mkHistory wQO) emptyLosers
  where
    go depth hist losers (deeds, state) = case terminate hist state of
        Continue hist' | sPECULATION -> continue depth hist' losers (deeds, state)
        _                            -> (losers, (deeds, state))
    
    continue depth hist losers (deeds, state@(Heap h _, _, _)) = (losers', (deeds'', (Heap (h'_winners' `M.union` h'_losers) ids'', k, in_e)))
      where
        (deeds', (Heap h' ids', k, in_e)) = reduce (deeds, state)
        
        -- It is *very important* that we prevent speculation from forcing heap-carried things that have proved
        -- to be losers in the past. This prevents us discovering speculation failure of some head binding, and
        -- then forcing several thunks that are each strict in it. This change was motivated by DigitsOfE2 -- it
        -- speeds up supercompilation of that benchmark massively.
        (h'_losers, h'_winners) | sPECULATE_ON_LOSERS = (M.empty, h')
                                | otherwise           = M.partition (\hb -> maybe False (`IS.member` losers) (heapBindingTag hb)) h'
        
        -- NB: It is important that we accumulate losers across "go" invocations in a state-monady kind of way, or DigitsOfE2 blows out
        -- even more than normal (it still takes 22s with this change).
        -- TODO: I suspect we should accumulate Losers across the boundary of speculate as well
        -- TODO: there is a difference between losers due to termination-halting and losers because we didn't have neough
        -- information available to complete evaluation
        (deeds'', Heap h'_winners' ids'', losers') = M.foldWithKey speculate_one (deeds', Heap h'_winners ids', losers) (h'_winners M.\\ h)
        speculate_one x' (Concrete in_e) (deeds, Heap h'_winners ids, losers)
{'fraction_non_alphanumeric': 0.10827178729689808, 'fraction_numerical': 0.0029542097488921715, 'mean_word_length': 2.2072477498815726, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 91, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4472397', 'n_tokens_mistral': 4141, 'n_tokens_neox': 4025, 'n_words': 1548}