TypeFlow-Based Monomorphism

Main.hs

module Main (main) where

import Data.HashMap.Strict qualified as H
import Data.Maybe (isJust)
import Data.Graph (stronglyConnComp, SCC(..), flattenSCC)
import Data.List (intercalate, nub)

-- See: https://dl.acm.org/doi/pdf/10.1145/3720472

data Type
    = Meta String
    | Basic String
    | Ctor TypeCtor
    deriving (Show, Eq)

data TypeCtor
    = Arrow Type Type
    | Record String [Type]
    deriving (Show, Eq)

data MetaFlowEdge
    = MetaFlowEdge
    { from :: String
    , target :: String
    , ctor :: Maybe TypeCtor
    }
    deriving (Show, Eq)

newtype MetaGraph = MetaGraph (H.HashMap String [MetaFlowEdge])
    deriving (Show, Eq)

data VisitState
    = Visiting {-# UNPACK #-} !Int
    | Done
    deriving (Show)

{- | Detects a cycle in the graph that contains at least one edge with a constructor.
If such a cycle exists, returns one edge with a constructor that is part of the cycle.

The algorithm uses Depth First Search (DFS) to traverse the graph.
It maintains the state of each node (Visiting or Done) to detect cycles.
Additionally, it tracks the "latest constructor edge" encountered on the current path.

When a back edge (u -> v) is found (indicating a cycle):
1. If the back edge itself has a constructor, we found a growing cycle. Return this edge.
2. If not, we check if there was any constructor edge on the path from v to u.
   We do this by comparing the depth of the latest constructor edge with the depth of v.
   If latestCtorDepth > depth(v), then the constructor edge is within the cycle. Return it.

Complexity: O(V + E) time and space, where V is vertices and E is edges.
-}
detectGrowingCycles :: MetaGraph -> Maybe MetaFlowEdge
detectGrowingCycles (MetaGraph graph) =
    let
        nodes = H.keys graph

        loop [] _ = Nothing
        loop (n : ns) visited
            | H.member n visited = loop ns visited
            | otherwise =
                case dfs n 0 0 Nothing visited of
                    (Just edge, _) -> Just edge
                    (Nothing, visited') -> loop ns visited'

        dfs :: String -> Int -> Int -> Maybe MetaFlowEdge -> H.HashMap String VisitState -> (Maybe MetaFlowEdge, H.HashMap String VisitState)
        dfs u depth latestCtorDepth latestCtorEdge visited =
            let
                visited' = H.insert u (Visiting depth) visited
                edges = H.lookupDefault [] u graph

                processEdges [] vMap = (Nothing, H.insert u Done vMap)
                processEdges (e : es) vMap =
                    let v = target e
                        hasCtor = isJust (ctor e)
                        newDepth = depth + 1
                        (newLatestDepth, newLatestEdge) =
                            if hasCtor
                                then (newDepth, Just e)
                                else (latestCtorDepth, latestCtorEdge)
                     in case H.lookup v vMap of
                            Just (Visiting vDepth) ->
                                if hasCtor
                                    then
                                        (Just e, vMap)
                                    else
                                        if newLatestDepth > vDepth
                                            then
                                                (newLatestEdge, vMap)
                                            else
                                                processEdges es vMap
                            Just Done ->
                                processEdges es vMap
                            Nothing ->
                                case dfs v newDepth newLatestDepth newLatestEdge vMap of
                                    (Just found, vMap') -> (Just found, vMap')
                                    (Nothing, vMap') -> processEdges es vMap'
             in
                processEdges edges visited'
     in
        loop nodes H.empty

newtype Solution = Solution (H.HashMap String [Type])
    deriving (Eq)

instance Show Solution where
    show (Solution sol) = unlines $ map showEntry $ H.toList sol
      where
        showEntry (k, ts) = k ++ ": " ++ showTypes ts
        showTypes [] = "[]"
        showTypes ts = "[" ++ intercalate ", " (map prettyType ts) ++ "]"

prettyType :: Type -> String
prettyType (Basic s) = s
prettyType (Meta s) = "?" ++ s
prettyType (Ctor (Arrow t1 t2)) = "(" ++ prettyType t1 ++ " -> " ++ prettyType t2 ++ ")"
prettyType (Ctor (Record name args)) = name ++ "<" ++ intercalate ", " (map prettyType args) ++ ">"

-- meta var's basic type mapping, this tells you the basic types that flow
-- into meta var.
newtype PrimitiveFlow = PrimitiveFlow (H.HashMap String [String])
    deriving (Show, Eq)

-- 1. no solution if there is a growing cycle
-- 2. otherwise, we can construct a solution by assigning basic types to given meta vars.
-- 3. Since there is no growing cycle, if a cycle ever exists, it just mean that
--    the meta vars in the SCC have the same solution.
-- 4. Hence we construct the topological order of SCCs and solve them in order.
-- 5. if x flows into y, then x's solutions are also y's solutions.
-- 6. the returned solution should map meta to all fully instantiated types.
--    that is, all metas in the body should be substituted with their solutions.
--    Multiple subsitutions are possible, we return all of them as this solution
--    is used to compute all types that out to be emitted for codegen.
solveMeta :: PrimitiveFlow -> MetaGraph -> Maybe Solution
solveMeta (PrimitiveFlow prims) g@(MetaGraph graph) =
    if isJust (detectGrowingCycles g)
        then Nothing
        else Just $ runSolve
  where
    -- Build adjacency list for SCC
    -- Nodes are all keys in graph and all targets
    allNodes = nub $ H.keys graph ++ concatMap (map target) (H.elems graph)
    
    -- Edge list for Data.Graph
    -- (node, key, [neighbors])
    -- We treat all edges as dependencies.
    adjList = map mkNode allNodes
      where
        mkNode u = (u, u, map target (H.lookupDefault [] u graph))

    sccs = stronglyConnComp adjList
    -- stronglyConnComp returns reverse topological order.
    -- We want topological order (upstream first).
    topo = reverse sccs

    -- Build reverse graph for easy lookup of incoming edges
    -- Map target -> [Edge]
    incomingEdges :: H.HashMap String [MetaFlowEdge]
    incomingEdges = H.fromListWith (++) $ do
        (_, edges) <- H.toList graph
        e <- edges
        return (target e, [e])

    runSolve :: Solution
    runSolve = Solution $ foldl processSCC initSol topo

    initSol :: H.HashMap String [Type]
    initSol = H.map (map Basic) prims

    processSCC :: H.HashMap String [Type] -> SCC String -> H.HashMap String [Type]
    processSCC currentSol comp =
        let nodes = flattenSCC comp
            -- Gather all types flowing into this SCC from outside
            -- or from PrimitiveFlow (already in currentSol)
            
            -- 1. Existing types from PrimitiveFlow for these nodes
            baseTypes = concatMap (\n -> H.lookupDefault [] n currentSol) nodes
            
            -- 2. Types from incoming edges
            flowTypes = concatMap (getIncomingTypes currentSol nodes) nodes
            
            allTypes = nub (baseTypes ++ flowTypes)
            
        in foldl (\s n -> H.insert n allTypes s) currentSol nodes

    getIncomingTypes :: H.HashMap String [Type] -> [String] -> String -> [Type]
    getIncomingTypes sol sccNodes node =
        let edges = H.lookupDefault [] node incomingEdges
        in concatMap (resolveEdge sol sccNodes) edges

    resolveEdge :: H.HashMap String [Type] -> [String] -> MetaFlowEdge -> [Type]
    resolveEdge sol sccNodes edge =
        let u = from edge
        in if u `elem` sccNodes
           then [] -- Ignore intra-SCC flow (handled by merging)
           else case ctor edge of
               Nothing -> H.lookupDefault [] u sol
               Just t  -> instantiate (Ctor t) sol

    instantiate :: Type -> H.HashMap String [Type] -> [Type]
    instantiate t sol =
        let vars = nub $ collectVars t
            assignments = generateAssignments vars sol
        in map (\env -> substitute env t) assignments

    collectVars :: Type -> [String]
    collectVars (Meta s) = [s]
    collectVars (Basic _) = []
    collectVars (Ctor (Arrow t1 t2)) = collectVars t1 ++ collectVars t2
    collectVars (Ctor (Record _ args)) = concatMap collectVars args

    generateAssignments :: [String] -> H.HashMap String [Type] -> [H.HashMap String Type]
    generateAssignments [] _ = [H.empty]
    generateAssignments (v:vs) sol =
        let rest = generateAssignments vs sol
            options = H.lookupDefault [] v sol
        in [H.insert v opt r | opt <- options, r <- rest]

    substitute :: H.HashMap String Type -> Type -> Type
    substitute env (Meta s) = H.lookupDefault (Meta s) s env
    substitute _ (Basic b) = Basic b
    substitute env (Ctor (Arrow t1 t2)) = Ctor (Arrow (substitute env t1) (substitute env t2))
    substitute env (Ctor (Record n args)) = Ctor (Record n (map (substitute env) args))

main :: IO ()
main = do
    putStrLn "=== Running detectGrowingCycles experiment ==="

    -- Example 1: A growing cycle
    -- A -> B (no ctor)
    -- B -> C (ctor: Record "R" [])
    -- C -> D (no ctor)
    -- D -> B (no ctor)
    -- Cycle: B -> C -> D -> B. Contains ctor edge B->C.
    -- G := { A -> B, R[B] -> C, C -> D, D -> B }
    let edgeAB = MetaFlowEdge "A" "B" Nothing
        edgeBC = MetaFlowEdge "B" "C" (Just (Record "R" [Meta "B"]))
        edgeCD = MetaFlowEdge "C" "D" Nothing
        edgeDB = MetaFlowEdge "D" "B" Nothing

        graph1 =
            MetaGraph $
                H.fromList
                    [ ("A", [edgeAB])
                    , ("B", [edgeBC])
                    , ("C", [edgeCD])
                    , ("D", [edgeDB])
                    ]

    putStrLn "Graph 1 (Growing Cycle):"
    print (detectGrowingCycles graph1)

    -- Example 2: A non-growing cycle
    -- X -> Y (ctor)
    -- Y -> Z (no ctor)
    -- Z -> Y (no ctor)
    -- Cycle: Y -> Z -> Y. No ctor in cycle. (X->Y is outside)
    -- G := { (Int -> X) -> Y, Y -> Z, Z -> Y }
    let edgeXY = MetaFlowEdge "X" "Y" (Just (Arrow (Basic "Int") (Meta "X")))
        edgeYZ = MetaFlowEdge "Y" "Z" Nothing
        edgeZY = MetaFlowEdge "Z" "Y" Nothing

        graph2 =
            MetaGraph $
                H.fromList
                    [ ("X", [edgeXY])
                    , ("Y", [edgeYZ])
                    , ("Z", [edgeZY])
                    ]

    putStrLn "\nGraph 2 (Non-Growing Cycle):"
    print (detectGrowingCycles graph2)

    -- Example 3: Complex growing cycle
    -- 1 -> 2 (no)
    -- 2 -> 3 (no)
    -- 3 -> 4 (ctor)
    -- 4 -> 2 (no)
    -- Cycle 2-3-4-2. Has ctor 3->4.
    let e12 = MetaFlowEdge "1" "2" Nothing
        e23 = MetaFlowEdge "2" "3" Nothing
        e34 = MetaFlowEdge "3" "4" (Just (Record "S" [Meta "3"]))
        e42 = MetaFlowEdge "4" "2" Nothing

        graph3 =
            MetaGraph $
                H.fromList
                    [ ("1", [e12])
                    , ("2", [e23])
                    , ("3", [e34])
                    , ("4", [e42])
                    ]

    putStrLn "\nGraph 3 (Complex Growing Cycle):"
    print (detectGrowingCycles graph3)

    putStrLn "\n=== Running solveMeta experiment ==="
    
    -- Example 4: Simple flow
    -- A -> B
    -- Prim: A = {Int}
    -- Result: A={Int}, B={Int}
    let eAB = MetaFlowEdge "A" "B" Nothing
        g4 = MetaGraph $ H.fromList [("A", [eAB])]
        p4 = PrimitiveFlow $ H.fromList [("A", ["Int"])]
    putStrLn "\nGraph 4 (Simple Flow A->B, A={Int}):"
    print (solveMeta p4 g4)

    -- Example 5: Ctor flow
    -- A -> B (ctor List<A>)
    -- Prim: A = {Int}
    -- Result: A={Int}, B={List<Int>}
    let eAB_ctor = MetaFlowEdge "A" "B" (Just (Record "List" [Meta "A"]))
        g5 = MetaGraph $ H.fromList [("A", [eAB_ctor])]
        p5 = PrimitiveFlow $ H.fromList [("A", ["Int"])]
    putStrLn "\nGraph 5 (Ctor Flow A->B, ctor=List<A>, A={Int}):"
    print (solveMeta p5 g5)

    -- Example 6: Cycle (SCC)
    -- A -> B, B -> A
    -- Prim: A = {Int}, B = {Bool}
    -- Result: A={Int, Bool}, B={Int, Bool}
    let eAB_cyc = MetaFlowEdge "A" "B" Nothing
        eBA_cyc = MetaFlowEdge "B" "A" Nothing
        g6 = MetaGraph $ H.fromList [("A", [eAB_cyc]), ("B", [eBA_cyc])]
        p6 = PrimitiveFlow $ H.fromList [("A", ["Int"]), ("B", ["Bool"])]
    putStrLn "\nGraph 6 (Cycle A<->B, A={Int}, B={Bool}):"
    print (solveMeta p6 g6)

    -- Example 7: Multiple substitutions
    -- A -> B (ctor Pair<A, A>)
    -- Prim: A = {Int, Bool}
    -- Result: B = {Pair<Int, Int>, Pair<Bool, Bool>}
    -- Note: Consistent instantiation means we don't get Pair<Int, Bool> etc.
    let eAB_pair = MetaFlowEdge "A" "B" (Just (Record "Pair" [Meta "A", Meta "A"]))
        g7 = MetaGraph $ H.fromList [("A", [eAB_pair])]
        p7 = PrimitiveFlow $ H.fromList [("A", ["Int", "Bool"])]
    putStrLn "\nGraph 7 (Multiple Subs A->B, ctor=Pair<A,A>, A={Int, Bool}):"
    print (solveMeta p7 g7)
    
    -- Example 8: Growing Cycle (Should fail)
    -- A -> A (ctor List<A>)
    let eAA = MetaFlowEdge "A" "A" (Just (Record "List" [Meta "A"]))
        g8 = MetaGraph $ H.fromList [("A", [eAA])]
        p8 = PrimitiveFlow $ H.fromList [("A", ["Int"])]
    putStrLn "\nGraph 8 (Growing Cycle A->A):"
    print (solveMeta p8 g8)

    -- Example 9: Consistent Instantiation & Multi-var
    -- A = {Int} (initially)
    -- B = {Bool}
    -- B -> A => A = {Int, Bool}
    -- A -> C (ctor Pair<A, B>)
    -- B -> C (ctor Pair<A, B>)
    -- Result C: Pair<Int, Bool>, Pair<Bool, Bool>
    -- (Consistent A, Consistent B)
    let eBA_9 = MetaFlowEdge "B" "A" Nothing
        eAC_9 = MetaFlowEdge "A" "C" (Just (Record "Pair" [Meta "A", Meta "B"]))
        eBC_9 = MetaFlowEdge "B" "C" (Just (Record "Pair" [Meta "A", Meta "B"]))
        
        g9 = MetaGraph $ H.fromList 
            [ ("B", [eBA_9, eBC_9])
            , ("A", [eAC_9])
            ]
        p9 = PrimitiveFlow $ H.fromList [("A", ["Int"]), ("B", ["Bool"])]
        
    putStrLn "\nGraph 9 (Consistent Pair<A, B>, A={Int, Bool}, B={Bool}):"
    print (solveMeta p9 g9)
    -- Example 10: Nested Pairs
    -- Extending Example 9 with D -> Pair<C, A>
    -- A = {Int, Bool}
    -- B = {Bool}
    -- C = {Pair<Int, Bool>, Pair<Bool, Bool>}
    -- D gets Pair<C, A>
    let eBA_10 = MetaFlowEdge "B" "A" Nothing
        eAC_10 = MetaFlowEdge "A" "C" (Just (Record "Pair" [Meta "A", Meta "B"]))
        eBC_10 = MetaFlowEdge "B" "C" (Just (Record "Pair" [Meta "A", Meta "B"]))
        eCD_10 = MetaFlowEdge "C" "D" (Just (Record "Pair" [Meta "C", Meta "A"]))
        eAD_10 = MetaFlowEdge "A" "D" (Just (Record "Pair" [Meta "C", Meta "A"]))
        
        g10 = MetaGraph $ H.fromList 
            [ ("B", [eBA_10, eBC_10])
            , ("A", [eAC_10, eAD_10])
            , ("C", [eCD_10])
            ]
        p10 = PrimitiveFlow $ H.fromList [("A", ["Int"]), ("B", ["Bool"])]
        
    putStrLn "\nGraph 10 (Nested Pair<C, A>):"
    print (solveMeta p10 g10)