<commit_before>
module Yhc.Core.Firstify.Super(super) where

import Yhc.Core hiding (uniqueBoundVarsCore, uniqueBoundVars)
import Yhc.Core.FreeVar3
import Yhc.Core.UniqueId

import Yhc.Core.Firstify.Mitchell.Util
import Yhc.Core.Firstify.Mitchell.Template
import Yhc.Core.Firstify.Mitchell.Terminate
import qualified Yhc.Core.Firstify.Mitchell.BiMap as BiMap

import Control.Exception
import Control.Monad
import Control.Monad.State
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.List
import Data.Maybe
import Debug.Trace
import Safe




type M a = State S a

data S = S {done :: Set.Set CoreFuncName -- those functions which have been done
           ,pending :: Set.Set CoreFuncName -- those which are being done
           ,core :: CoreFuncMap -- the entire program
           
           ,special :: BiMap.BiMap CoreFuncName CoreExpr -- which special variants do we have
           ,terminate :: () -- termination check
           
           ,varId :: Int -- what is the next variable id to use
           ,funcId :: Int -- what is the next function id to use
           }


instance UniqueId S where
    getId = varId
    putId x s = s{varId = x}


super :: Core -> Core
super c = fromCoreFuncMap c $ core $ flip execState undefined $ do
        c <- return $ ensureInvariants [NoRecursiveLet,NoCorePos] c
        let s0 = S Set.empty Set.empty undefined BiMap.empty () 0 (uniqueFuncsNext c)
        c <- uniqueBoundVarsCore c
        modify $ \s -> s{core = toCoreFuncMap c}
        foFunc "main"


foFunc :: CoreFuncName -> M CoreExpr
foFunc x = do
    s <- get
    func <- return $ coreFuncMap (core s) x
    when (isCoreFunc func && x `Set.notMember` done s && x `Set.notMember` pending s) $ do
        modify $ \s -> s{pending = Set.insert x (pending s)}
        (args,body) <- liftM fromCoreLam $ foExpr (coreFuncBody func)
        modify $ \s -> s{core = Map.insert x (CoreFunc x (coreFuncArgs func ++ args) body) (core s)
                        ,pending = Set.delete x (pending s)
                        ,done = Set.insert x (done s)
                        }

    func <- return $ coreFuncMap (core s) x
    vars <- getVars (coreFuncArity func)
    return $ CoreLam vars  $CoreApp (CoreFun x) (map CoreVar vars)



foExpr :: CoreExpr -> M CoreExpr
foExpr = transformM fo


fo :: CoreExpr -> M CoreExpr
fo x = return x

{-




-- In each step first inline all top-level function bindings
-- and let's that appear to be bound to an unsaturated
--
-- Then specialise each value
step :: CoreFuncMap -> SS CoreFuncMap
step = f acts
    where
        (*) = (,)
        acts = ["lambdas" * lambdas, "simplify" * simplify, "inline" * inline, "specialise" * specialise]<commit_msg>Add loads to Super, start of an initial implementation<commit_after>
module Yhc.Core.Firstify.Super(super) where

import Yhc.Core hiding (uniqueBoundVarsCore, uniqueBoundVars)
import Yhc.Core.FreeVar3
import Yhc.Core.UniqueId

import Yhc.Core.Firstify.Mitchell.Util
import Yhc.Core.Firstify.Mitchell.Template
import Yhc.Core.Firstify.Mitchell.Terminate
import qualified Yhc.Core.Firstify.Mitchell.BiMap as BiMap

import Control.Exception
import Control.Monad
import Control.Monad.State
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.List
import Data.Maybe
import Debug.Trace
import Safe




type M a = State S a

data S = S {done :: Set.Set CoreFuncName -- those functions which have been done
           ,pending :: Set.Set CoreFuncName -- those which are being done
           ,core :: CoreFuncMap -- the entire program
           
           ,special :: BiMap.BiMap CoreFuncName CoreExpr -- which special variants do we have
           ,terminate :: () -- termination check
           
           ,varId :: Int -- what is the next variable id to use
           ,funcId :: Int -- what is the next function id to use
           }


instance UniqueId S where
    getId = varId
    putId x s = s{varId = x}


super :: Core -> Core
super c = fromCoreFuncMap c $ core $ flip execState undefined $ do
        c <- return $ ensureInvariants [NoRecursiveLet,NoCorePos] c
        let s0 = S Set.empty Set.empty undefined BiMap.empty () 0 (uniqueFuncsNext c)
        put (s0 :: S)
        c <- uniqueBoundVarsCore c
        modify $ \s -> s{core = toCoreFuncMap c}
        foFunc "main"


foFunc :: CoreFuncName -> M Int
foFunc x = do
    s <- get
    func <- return $ coreFuncMap (core s) x
    when (isCoreFunc func && x `Set.notMember` done s && x `Set.notMember` pending s) $ do
        modify $ \s -> s{pending = Set.insert x (pending s)}
        (args,body) <- liftM fromCoreLam $ fo (coreFuncBody func)
        modify $ \s -> s{core = Map.insert x (CoreFunc x (coreFuncArgs func ++ args) body) (core s)
                        ,pending = Set.delete x (pending s)
                        ,done = Set.insert x (done s)
                        }
    return $ coreFuncArity $ coreFuncMap (core s) x



-- go downwards
-- CoreApp (CoreFun _) _
--     first saturate the call with its arity
--     then look towards templates


fo :: CoreExpr -> M CoreExpr
fo (CoreFun x) = fo (CoreApp (CoreFun x) [])

fo (CoreApp (CoreLam vs x) xs) = fo $
        CoreApp (CoreLam vs2 x2) xs2
    where
        x2 = coreLet bind $ replaceFreeVars rep x

        n = min (length vs) (length xs)
        (vs1,vs2) = splitAt n vs
        (xs1,xs2) = splitAt n xs

        (rep,bind) = partition isOne (zip vs1 xs1)
        isOne (v,_) = countFreeVar v x <= 1


fo (CoreApp (CoreFun x) xs) = do
    arity <- foFunc x
    vs <- getVars $ max 0 (arity - length xs)
    xs <- return $ xs ++ map CoreVar vs
    o <- return $ CoreApp (CoreFun x) xs

    s <- get
    let t = templateCreate (isCorePrim . coreFuncMap (core s)) o
    o <- if t == templateNone then return o else do
        let tfull = templateExpand (`BiMap.lookup` special s) t
            holes = templateHoles o t
        case BiMap.lookupRev t (special s) of
            -- OPTION 1: Not previously done, and a homeomorphic embedding
            --Nothing | not $ askSpec within tfull (terminate s) -> return x
            -- OPTION 2: Previously done
            Just name ->
                return $ coreApp (CoreFun name) holes
            -- OPTION 3: New todo
            done -> do
                let name = uniqueJoin (templateName t) (funcId s)
                fun <- templateGenerate (coreFuncMap (core s)) name t
                modify $ \s -> s
                    { {-terminate = addSpec name tfull $
                                   cloneSpec within name $ terminate s
                      , -} funcId = funcId s + 1
                    ,special = BiMap.insert name t (special s)
                    ,core = Map.insert name fun (core s)
                    }
                return $ coreApp (CoreFun name) holes
    return $ coreLam vs o

fo x = descendM fo x

{-




-- In each step first inline all top-level function bindings
-- and let's that appear to be bound to an unsaturated
--
-- Then specialise each value
step :: CoreFuncMap -> SS CoreFuncMap
step = f acts
    where
        (*) = (,)
        acts = ["lambdas" * lambdas, "simplify" * simplify, "inline" * inline, "specialise" * specialise]
{'fraction_non_alphanumeric': 0.09293938132889444, 'fraction_numerical': 0.003190456373976973, 'mean_word_length': 2.1595092024539877, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '742355', 'n_tokens_mistral': 2212, 'n_tokens_neox': 2101, 'n_words': 802}