<commit_before>-- | examine all uses of types in a program to determine which ones are
-- actually needed in the method generation

module E.TypeAnalysis(typeAnalyze, Typ(), pruneE) where

import Control.Monad.Reader
import Control.Monad.Identity
import Data.Monoid
import Data.FunctorM
import qualified Data.Set as Set
import qualified Data.Map as Map

import E.Annotate
import E.E hiding(isBottom)
import E.Inline(emapE',emapE_)
import E.Program
import E.TypeCheck
import Fixer.Fixer
import Fixer.VMap
import GenUtil
import Info.Types
import Name.Name
import Name.Names
import qualified Info.Info as Info
import Support.CanType


type Typ = VMap Name
type Env = (Value (Set.Set TVr),Map.Map Id [Value Typ])

extractValMap :: [(TVr,E)] -> Map.Map Id [Value Typ]
extractValMap ds = Map.fromList [ (tvrIdent t,f e []) | (t,e) <- ds] where
    f (ELam tvr e) rs | sortStarLike (getType tvr) = f e (runIdentity (Info.lookup $ tvrInfo tvr):rs)
    f _ rs = reverse rs

-- all variables _must_ be unique before running this
typeAnalyze :: Program -> IO Program
typeAnalyze prog = do
    fixer <- newFixer
    usedVals <- newValue fixer Set.empty
    let lambind _ nfo = do
            x <- newValue fixer ( bottom :: Typ)
            return $ Info.insert x nfo
        lamread _ nfo | Just v <- Info.lookup nfo = do
            rv <- readValue v
            return (Info.insert (rv :: Typ) $ Info.delete (undefined :: Value Typ) nfo)
        lamread _ nfo = return nfo
        lamdel _ nfo = return (Info.delete (undefined :: Value Typ) nfo)
    prog <- annotateProgram mempty lambind (\_ -> return) (\_ -> return) prog
    let ds = programDs prog
        env = (usedVals,extractValMap ds)
        entries = progEntryPoints prog ++ [ t | (t,_) <- ds, getProperty prop_INSTANCE t]
    calcDs env ds
    mapM_ (calcE env . EVar ) entries
    mapM_ (sillyEntry env) entries
    calcFixpoint "type analysis" fixer
    prog <- annotateProgram mempty (\_ -> return) (\_ -> return) lamread prog
    prog <- annotateProgram mempty lamdel (\_ -> return) (\_ -> return) prog
    return prog

sillyEntry :: Env -> TVr -> IO ()
sillyEntry env t = mapM_ (addRule . (`isSuperSetOf` value (vmapSingleton v_silly))) args where
    args = lookupArgs t env

lookupArgs t (_,tm) = maybe [] id (Map.lookup (tvrIdent t) tm)

calcDs ::  Env -> [(TVr,E)] -> IO ()
calcDs env@(usedVals,_) ds = do
    mapM_ d ds
    flip mapM_ ds $ \ (v,e) -> do
        addRule $ conditionalRule (v `Set.member`) usedVals (ioToRule $ calcE env e)
     where
        --mapM_ d ds >> mapM_ (calcE env) (snds ds) where
    d (t,e) | not (sortStarLike (getType t)) = return ()
    d (t,e) | Just v <- getValue e = do
        let Just t' = Info.lookup (tvrInfo t)
        addRule $ t' `isSuperSetOf` v
    d (t, ELit (LitCons n xs _)) = do
        let Just t' = Info.lookup (tvrInfo t)
            v = vmapSingleton n
        addRule $ t' `isSuperSetOf` (value v)
        xs' <- mapM getValue xs
        flip mapM_ (zip xs' [0.. ])  $ \ (v,i) -> do
            addRule $ modifiedSuperSetOf t' v (vmapArgSingleton n i)
    d (t, EPi TVr { tvrType = a} b) = do
        let Just t' = Info.lookup (tvrInfo t)
            v = vmapSingleton tc_Arrow
        addRule $ t' `isSuperSetOf` (value v)
        xs' <- mapM getValue [a,b]
        flip mapM_ (zip xs' [0.. ])  $ \ (v,i) -> do
            addRule $ modifiedSuperSetOf t' v (vmapArgSingleton tc_Arrow i)
    d (t,e) | (EVar v,as) <- fromAp e = do
        let Just t' = Info.lookup (tvrInfo t)
            Just v' = Info.lookup (tvrInfo v)
        as' <- mapM getValue as
        addRule $ dynamicRule v' $ \ v -> mconcat $ flip map (vmapHeads v) $ \ h ->
            mconcat $ t' `isSuperSetOf` value (vmapSingleton h) : (flip map (zip as' [0.. ])  $ \ (a,i) -> modifiedSuperSetOf t' a $ \ v -> vmapArgSingleton h i v)
    d (t,e) = fail $ "calcDs: " ++ show (t,e)

-- TODO - make default case conditional
calcAlt env v (Alt (LitCons n xs _) e) = do
    addRule $ conditionalRule (\ (VMap _ vs) -> n `Set.member` vs) v $ ioToRule $ do
        calcE env e
        flip mapM_ (zip [0..] xs) $ \ (i,t) -> do
            let Just t' = Info.lookup (tvrInfo t)
            addRule $ modifiedSuperSetOf t' v (vmapArg n i)


calcE :: Env -> E -> IO ()
calcE (usedVals,env) (ELetRec ds e) = calcDs nenv ds >> calcE nenv e where
    nenv = (usedVals,extractValMap ds `Map.union` env)
calcE env e | (e',(_:_)) <- fromLam e = calcE env e'
calcE env ec@ECase {} | sortStarLike (getType $ eCaseScrutinee ec) = do
    calcE env (eCaseScrutinee ec)
    fmapM_ (calcE env) (eCaseDefault ec)
    v <- getValue (eCaseScrutinee ec)
    mapM_ (calcAlt env v) (eCaseAlts ec)
calcE env ec@ECase {} = do
    calcE env (eCaseScrutinee ec)
    mapM_ (calcE env) (caseBodies ec)
calcE env e@ELit {} = tagE env e
calcE env e@EPrim {} = tagE env e
calcE _ EError {} = return ()
calcE _ ESort {} = return ()
calcE _ Unknown = return ()
calcE env e | (EVar v,as@(_:_)) <- fromAp e, Just ts <- Map.lookup (tvrIdent v) (snd env) = do
    tagE env e
    when (length as < length ts) $ fail "calcE: unsaturated call to function"
    flip mapM_ (zip as ts) $ \ (a,t) -> do
        when (sortStarLike (getType a)) $ do
            a' <- getValue a
            addRule $ t `isSuperSetOf` a'
calcE env e@EVar {} = tagE env e
calcE env e@EAp {} = tagE env e
calcE _ e = fail $ "odd calcE: " ++ show e

tagE (usedVals,_) (EVar v) = addRule $ usedVals `isSuperSetOf` value (Set.singleton v)
tagE env e  = emapE_ (tagE env) e

getValue (EVar v)
    | Just x <- Info.lookup (tvrInfo v) = return x
    | otherwise = fail $ "getValue: no varinfo: " ++ show v
getValue e | Just c <- typConstant e = return $ value c
getValue e = fail $ "getValue: " ++ show e

typConstant :: Monad m => E -> m Typ
typConstant (EPi TVr { tvrType = a} b) = do
    ab <- mapM typConstant [a,b]
    return $ vmapValue tc_Arrow ab
typConstant (ELit (LitCons n xs _)) = do
    xs' <- mapM typConstant xs
    return $ vmapValue n xs'
typConstant e = fail $ "typConstant: " ++ show e


-- pruning the unused branches of typecase statements


pruneE :: E -> IO E
pruneE e = return $ runIdentity (prune e)  where
    prune ec@ECase { eCaseScrutinee = EVar v } | sortStarLike (getType v), Just (VMap _ ns) <- Info.lookup (tvrInfo v) = do
        ec' <- pruneCase ec ns
        emapE' prune ec'
    prune e = emapE' prune e

pruneCase :: Monad m => E -> Set.Set Name -> m E
pruneCase ec ns = return $ if null (caseBodies nec) then err else nec where
    err = EError "pruneCase: all alternatives pruned" (getType ec)
    nec = ec { eCaseAlts = f [] $ eCaseAlts ec, eCaseDefault = cd (eCaseDefault ec)}
    f xs [] = reverse xs
    f xs (alt@(Alt (LitCons n _ _) _):rs) | not (n `Set.member` ns) = f xs rs
    f xs (alt:rs) = f (alt:xs) rs
    cd (Just d) | or [ n `notElem` as | n <- Set.toList ns ] = Just d
    cd Nothing = Nothing
    -- The reason we do this is because for a typecase, we need a valid default in order to get the most general type
    cd (Just d) = Just $ EError "pruneCase: default pruned" (getType d)
    as = [ n | LitCons n _ _ <- casePats ec ]


{-

specializeProgram :: Program -> Program
specializeProgram prog = ans where
    entries = Set.fromList $ progEntryPoints -- must not be specialized
    ans = runReader (programMapDs f prog) (cenv $ programDs prog)
    f (t,e) = do
        env <- ask
        ne <- case Map.lookup t env of
            Just (cd,_) -> cd e
            Nothing -> return e
        ne' <- de e
        return (t,ne')
    cenv ds = undefined
    de = undefined

-}<commit_msg>utilize new proxy support in Fixer.VMap<commit_after>-- | examine all uses of types in a program to determine which ones are
-- actually needed in the method generation

module E.TypeAnalysis(typeAnalyze, Typ(), pruneE) where

import Control.Monad.Reader
import Control.Monad.Identity
import Data.Monoid
import Data.FunctorM
import qualified Data.Set as Set
import qualified Data.Map as Map

import E.Annotate
import E.E hiding(isBottom)
import E.Inline(emapE',emapE_)
import E.Program
import E.TypeCheck
import Fixer.Fixer
import Fixer.VMap
import GenUtil
import Info.Types
import Name.Name
import Name.Names
import qualified Info.Info as Info
import Support.CanType


type Typ = VMap () Name
type Env = (Value (Set.Set TVr),Map.Map Id [Value Typ])

extractValMap :: [(TVr,E)] -> Map.Map Id [Value Typ]
extractValMap ds = Map.fromList [ (tvrIdent t,f e []) | (t,e) <- ds] where
    f (ELam tvr e) rs | sortStarLike (getType tvr) = f e (runIdentity (Info.lookup $ tvrInfo tvr):rs)
    f _ rs = reverse rs

-- all variables _must_ be unique before running this
typeAnalyze :: Program -> IO Program
typeAnalyze prog = do
    fixer <- newFixer
    usedVals <- newValue fixer Set.empty
    let lambind _ nfo = do
            x <- newValue fixer ( bottom :: Typ)
            return $ Info.insert x nfo
        lamread _ nfo | Just v <- Info.lookup nfo = do
            rv <- readValue v
            return (Info.insert (rv :: Typ) $ Info.delete (undefined :: Value Typ) nfo)
        lamread _ nfo = return nfo
        lamdel _ nfo = return (Info.delete (undefined :: Value Typ) nfo)
    prog <- annotateProgram mempty lambind (\_ -> return) (\_ -> return) prog
    let ds = programDs prog
        env = (usedVals,extractValMap ds)
        entries = progEntryPoints prog ++ [ t | (t,_) <- ds, getProperty prop_INSTANCE t]
    calcDs env ds
    mapM_ (calcE env . EVar ) entries
    mapM_ (sillyEntry env) entries
    calcFixpoint "type analysis" fixer
    prog <- annotateProgram mempty (\_ -> return) (\_ -> return) lamread prog
    prog <- annotateProgram mempty lamdel (\_ -> return) (\_ -> return) prog
    return prog

sillyEntry :: Env -> TVr -> IO ()
sillyEntry env t = mapM_ (addRule . (`isSuperSetOf` value (vmapPlaceholder ()))) args where
    args = lookupArgs t env

lookupArgs t (_,tm) = maybe [] id (Map.lookup (tvrIdent t) tm)

calcDs ::  Env -> [(TVr,E)] -> IO ()
calcDs env@(usedVals,_) ds = do
    mapM_ d ds
    flip mapM_ ds $ \ (v,e) -> do
        addRule $ conditionalRule (v `Set.member`) usedVals (ioToRule $ calcE env e)
     where
        --mapM_ d ds >> mapM_ (calcE env) (snds ds) where
    d (t,e) | not (sortStarLike (getType t)) = return ()
    d (t,e) | Just v <- getValue e = do
        let Just t' = Info.lookup (tvrInfo t)
        addRule $ t' `isSuperSetOf` v
    d (t, ELit (LitCons n xs _)) = do
        let Just t' = Info.lookup (tvrInfo t)
            v = vmapSingleton n
        addRule $ t' `isSuperSetOf` (value v)
        xs' <- mapM getValue xs
        flip mapM_ (zip xs' [0.. ])  $ \ (v,i) -> do
            addRule $ modifiedSuperSetOf t' v (vmapArgSingleton n i)
    d (t, EPi TVr { tvrType = a} b) = do
        let Just t' = Info.lookup (tvrInfo t)
            v = vmapSingleton tc_Arrow
        addRule $ t' `isSuperSetOf` (value v)
        xs' <- mapM getValue [a,b]
        flip mapM_ (zip xs' [0.. ])  $ \ (v,i) -> do
            addRule $ modifiedSuperSetOf t' v (vmapArgSingleton tc_Arrow i)
    d (t,e) | (EVar v,as) <- fromAp e = do
        let Just t' = Info.lookup (tvrInfo t)
            Just v' = Info.lookup (tvrInfo v)
        as' <- mapM getValue as
        addRule $ dynamicRule v' $ \ v -> mconcat $ (t' `isSuperSetOf` value (vmapDropArgs v)):case vmapHeads v of
                Just vs -> concat $ flip map vs $ \h -> (flip map (zip as' [0.. ])  $ \ (a,i) -> modifiedSuperSetOf t' a $ \ v -> vmapArgSingleton h i v)
                Nothing -> flip map (zip as' [0.. ])  $ \ (_,i) -> isSuperSetOf t' (value $ vmapProxyIndirect i v)
    d (t,e) = fail $ "calcDs: " ++ show (t,e)

-- TODO - make default case conditional
calcAlt env v (Alt (LitCons n xs _) e) = do
    addRule $ conditionalRule (n `vmapMember`) v $ ioToRule $ do
        calcE env e
        flip mapM_ (zip [0..] xs) $ \ (i,t) -> do
            let Just t' = Info.lookup (tvrInfo t)
            addRule $ modifiedSuperSetOf t' v (vmapArg n i)


calcE :: Env -> E -> IO ()
calcE (usedVals,env) (ELetRec ds e) = calcDs nenv ds >> calcE nenv e where
    nenv = (usedVals,extractValMap ds `Map.union` env)
calcE env e | (e',(_:_)) <- fromLam e = calcE env e'
calcE env ec@ECase {} | sortStarLike (getType $ eCaseScrutinee ec) = do
    calcE env (eCaseScrutinee ec)
    fmapM_ (calcE env) (eCaseDefault ec)
    v <- getValue (eCaseScrutinee ec)
    mapM_ (calcAlt env v) (eCaseAlts ec)
calcE env ec@ECase {} = do
    calcE env (eCaseScrutinee ec)
    mapM_ (calcE env) (caseBodies ec)
calcE env e@ELit {} = tagE env e
calcE env e@EPrim {} = tagE env e
calcE _ EError {} = return ()
calcE _ ESort {} = return ()
calcE _ Unknown = return ()
calcE env e | (EVar v,as@(_:_)) <- fromAp e, Just ts <- Map.lookup (tvrIdent v) (snd env) = do
    tagE env e
    when (length as < length ts) $ fail "calcE: unsaturated call to function"
    flip mapM_ (zip as ts) $ \ (a,t) -> do
        when (sortStarLike (getType a)) $ do
            a' <- getValue a
            addRule $ t `isSuperSetOf` a'
calcE env e@EVar {} = tagE env e
calcE env e@EAp {} = tagE env e
calcE _ e = fail $ "odd calcE: " ++ show e

tagE (usedVals,_) (EVar v) = addRule $ usedVals `isSuperSetOf` value (Set.singleton v)
tagE env e  = emapE_ (tagE env) e

getValue (EVar v)
    | Just x <- Info.lookup (tvrInfo v) = return x
    | otherwise = fail $ "getValue: no varinfo: " ++ show v
getValue e | Just c <- typConstant e = return $ value c
getValue e = fail $ "getValue: " ++ show e

typConstant :: Monad m => E -> m Typ
typConstant (EPi TVr { tvrType = a} b) = do
    ab <- mapM typConstant [a,b]
    return $ vmapValue tc_Arrow ab
typConstant (ELit (LitCons n xs _)) = do
    xs' <- mapM typConstant xs
    return $ vmapValue n xs'
typConstant e = fail $ "typConstant: " ++ show e


-- pruning the unused branches of typecase statements


pruneE :: E -> IO E
pruneE e = return $ runIdentity (prune e)  where
    prune ec@ECase { eCaseScrutinee = EVar v } | sortStarLike (getType v), Just vm <- Info.lookup (tvrInfo v) = do
        ec' <- pruneCase ec vm
        emapE' prune ec'
    prune e = emapE' prune e

pruneCase :: (Monad m) => E -> VMap () Name -> m E
pruneCase ec ns = return $ if null (caseBodies nec) then err else nec where
    err = EError "pruneCase: all alternatives pruned" (getType ec)
    nec = ec { eCaseAlts = f [] $ eCaseAlts ec, eCaseDefault = cd (eCaseDefault ec)}
    f xs [] = reverse xs
    f xs (alt@(Alt (LitCons n _ _) _):rs) | not (n `vmapMember` ns) = f xs rs
    f xs (alt:rs) = f (alt:xs) rs
    cd (Just d) | Just nns <- vmapHeads ns, or [ n `notElem` as | n <- nns ] = Just d
                | Nothing <- vmapHeads ns = Just d
    cd Nothing = Nothing
    -- The reason we do this is because for a typecase, we need a valid default in order to get the most general type
    cd (Just d) = Just $ EError "pruneCase: default pruned" (getType d)
    as = [ n | LitCons n _ _ <- casePats ec ]


{-

specializeProgram :: Program -> Program
specializeProgram prog = ans where
    entries = Set.fromList $ progEntryPoints -- must not be specialized
    ans = runReader (programMapDs f prog) (cenv $ programDs prog)
    f (t,e) = do
        env <- ask
        ne <- case Map.lookup t env of
            Just (cd,_) -> cd e
            Nothing -> return e
        ne' <- de e
        return (t,ne')
    cenv ds = undefined
    de = undefined

-}
{'fraction_non_alphanumeric': 0.11908805853148681, 'fraction_numerical': 0.0005879278808466162, 'mean_word_length': 2.5355658198614317, 'pattern_counts': {'":': 0, '<': 69, '<?xml version=': 0, '>': 79, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3050560', 'n_tokens_mistral': 5514, 'n_tokens_neox': 5309, 'n_words': 1961}