<commit_before>module Language.Haskell.Liquid.Synthesis where

import qualified Data.Graph as G
import           Data.Hashable
import           Data.HashSet (HashSet)
import qualified Data.HashSet as Set
import qualified Data.HashMap.Strict as M
import           Prelude

import qualified Language.Fixpoint.Types                    as F
import           Language.Haskell.Liquid.Types
import           Util (sndOf3)
import           Var


-- JP: What if there are multiple holes in the expression?
-- 
-- synthesize :: Var -> SpecType -> Env -> [Expr]
-- synthesize _ _ _ = []

-- data Graph a = Graph {
--       rootNodes :: HashSet a
--     , graphEdges :: M.HashMap a (HashSet a)
--     }

-- | Builds a dependency graph where an edge from hole a to hole b indicates hole b is dependent on hole a.
-- Root holes are not dependent on other holes.
-- Returns Nothing if a cycle is detected.
holeDependencyOrder :: M.HashMap Var (HoleInfo SpecType) -> Maybe [Var]
holeDependencyOrder holeMap = 
    let seen = Set.empty in
    let holes = M.keys holeMap in

    -- Find hole's dependencies.
    let deps = map (zipL $ findDependencies seen) holes in

    -- Return Nothing if cycle found.
    case sequence deps of
        Nothing ->
            Nothing

        Just deps ->

            -- Build graph.
            let (graph, nodeFromVertex, vertexFromKey) = G.graphFromEdges deps in

            -- Sort by partial ordering.
            let vertices = reverse $ G.topSort graph in -- JP: Use revTopSort if they merge.

            Just $ map (sndOf3 . nodeFromVertex) vertices


            -- -- Invert edges.
            -- let edges = M.fromListWith mappend $ invert deps in

            -- -- Find root nodes.
            -- let roots = findRoots edges in

            -- -- Build graph.
            -- Just $ Graph roots edges


    where
        zipL :: (a -> Maybe [a]) -> a -> Maybe ((), a, [a])
        zipL f x = (\y -> ((),x,y)) <$> f x

        -- Find all holes that this hole is dependent on.
        findDependencies :: HashSet Var -> Var -> Maybe [Var]
        findDependencies seen v | v `Set.member` seen = 
            if v `M.member` holeMap then 
                -- If we've seen this hole before, there's a cycle.
                Nothing
            else
                -- Otherwise, we're done.
                Just []
                
        findDependencies seen' v =
            -- Add v to seen.
            let seen = Set.insert v seen' in

            -- Lookup refinement type of v in environment.
            let reft = lookupAREnv v in

            -- Get set of variables in reft of v.
            -- Remove v from set of variables.
            -- Recursively find dependencies of set of variables.
            -- Check if there are any cycles in dependencies.
            -- Union all dependencies.
            -- Return all dependencies that are holes.
            
            error "TODO"

        -- JP: Does a function like this already exist?
        -- What's the right way to convert Var to Symbol?
        lookupAREnv :: Var -> AREnv t -> Maybe t
        lookupAREnv v REnv{..} = 
            let s = F.symbol v in

            -- Lookup local first.
            case M.lookup s reLocal of
                Just r -> 
                    Just r
                Nothing -> 
                    -- Lookup global next.
                    M.lookup s reGlobal
                        
        -- -- Expand edges and reverse the direction.
        -- invert :: Hashable a => [(a,[a])] -> [(a,HashSet a)]
        -- invert = concatMap (\(a, bs) -> map (\b -> ( b, Set.singleton a)) bs)
        -- 
        -- -- Find root nodes (Nodes that do not have edges pointed towards them).
        -- findRoots :: (Hashable a, Eq a) => M.HashMap a (HashSet a) -> HashSet a
        -- findRoots m =
        --     let holes = M.keysSet m in
        --     M.foldr (\towards acc -> Set.difference acc towards) holes m


        


<commit_msg>return list of pairs instead<commit_after>module Language.Haskell.Liquid.Synthesis where

import qualified Data.Graph as G
import           Data.Hashable
import           Data.HashSet (HashSet)
import qualified Data.HashSet as Set
import qualified Data.HashMap.Strict as M
import           Prelude

import qualified Language.Fixpoint.Types                    as F
import           Language.Haskell.Liquid.Types
import           Util (fstOf3)
import           Var


-- | Builds a dependency graph where an edge from hole a to hole b indicates hole b is dependent on hole a.
-- Root holes are not dependent on other holes.
-- Returns Nothing if a cycle is detected.
holeDependencyOrder :: M.HashMap Var (HoleInfo SpecType) -> Maybe [(Var, HoleInfo SpecType)]
holeDependencyOrder holeMap = 
    let seen = Set.empty in
    let holes = M.toList holeMap in

    -- Find hole's dependencies.
    let deps = map (zipL $ findDependencies seen) holes in

    -- Return Nothing if cycle found.
    case sequence deps of
        Nothing ->
            Nothing

        Just deps ->

            -- Build graph.
            let (graph, nodeFromVertex, _vertexFromKey) = G.graphFromEdges deps in

            -- Sort by partial ordering.
            let vertices = reverse $ G.topSort graph in -- JP: Use revTopSort if they merge.

            Just $ map (fstOf3 . nodeFromVertex) vertices


    where
        zipL :: ((k,v) -> Maybe [k]) -> (k,v) -> Maybe ((k,v), k, [k])
        zipL f x@(k,_) = (\y -> (x,k,y)) <$> f x

        -- Find all holes that this hole is dependent on.
        findDependencies :: HashSet Var -> (Var, HoleInfo SpecType) -> Maybe [Var]
        findDependencies seen (v, _) | v `Set.member` seen = 
            if v `M.member` holeMap then 
                -- If we've seen this hole before, there's a cycle.
                Nothing
            else
                -- Otherwise, we're done.
                Just [] -- JP: Should we return Nothing here too?
                
        findDependencies seen' (v, hi) =
            -- Add v to seen.
            let seen = Set.insert v seen' in

            -- Lookup refinement type of v in environment.
            let reft = lookupAREnv v $ henv hi in

            -- Get set of variables in reft of v.
            -- Remove v from set of variables.
            -- Recursively find dependencies of set of variables.
            -- Check if there are any cycles in dependencies.
            -- Union all dependencies.
            -- Return all dependencies that are holes.
            
            error "TODO"

        -- JP: Does a function like this already exist?
        -- What's the right way to convert Var to Symbol?
        lookupAREnv :: Var -> AREnv t -> Maybe t
        lookupAREnv v REnv{..} = 
            let s = F.symbol v in

            -- Lookup local first.
            case M.lookup s reLocal of
                Just r -> 
                    Just r
                Nothing -> 
                    -- Lookup global next.
                    M.lookup s reGlobal
                        
{'fraction_non_alphanumeric': 0.08633910820789549, 'fraction_numerical': 0.0005680204487361545, 'mean_word_length': 1.481677237491191, 'pattern_counts': {'":': 0, '<': 5, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4656387', 'n_tokens_mistral': 1910, 'n_tokens_neox': 1796, 'n_words': 764}