module Unison.Util.Components where

import Data.Graph qualified as Graph
import Data.Map qualified as Map
import Data.Set qualified as Set
import Unison.Prelude

-- | Order bindings by dependencies and group into components.
-- Each component consists of > 1 bindings, each of which depends
-- transitively on all other bindings in the component.
--
-- 1-element components may or may not depend on themselves.
--
-- The order is such that a component at index i will not depend
-- on components and indexes > i. But a component at index i does not
-- _necessarily_ depend on any components at earlier indices.
--
-- Example:
--
--   let rec
--     ping n = pong (n + 1);
--     pong n = ping (n + 1);
--     g = id 42;
--     y = id "hi"
--     id x = x;
--   in ping g
--
-- `components` would produce `[[ping,pong], [id], [g], [y]]`
-- Notice that `id` comes before `g` and `y` in the output, since
-- both `g` and `y` depend on `id`.
--
-- Uses Tarjan's algorithm:
--   https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm
components :: (Ord v) => (t -> Set v) -> [(v, t)] -> [[(v, t)]]
components :: forall v t. Ord v => (t -> Set v) -> [(v, t)] -> [[(v, t)]]
components t -> Set v
freeVars [(v, t)]
bs =
  let varIds :: Map v Int
varIds =
        [(v, Int)] -> Map v Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (((v, t) -> v) -> [(v, t)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (v, t) -> v
forall a b. (a, b) -> a
fst [(v, t)]
bs [v] -> [Int] -> [(v, Int)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Int] -> [Int]
forall a. [a] -> [a]
reverse [(Int
1 :: Int) .. [(v, t)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(v, t)]
bs])
      -- something horribly wrong if this bombs
      msg :: a
msg = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Components.components bug"
      varId :: v -> Int
varId v
v = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
forall {a}. a
msg (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ v -> Map v Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup v
v Map v Int
varIds

      -- use ints as keys for graph to preserve original source order as much as
      -- possible
      graph :: [((v, t), Int, [Int])]
graph = [((v
v, t
b), v -> Int
varId v
v, t -> [Int]
deps t
b) | (v
v, t
b) <- [(v, t)]
bs]
      vars :: Set v
vars = [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList (((v, t) -> v) -> [(v, t)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (v, t) -> v
forall a b. (a, b) -> a
fst [(v, t)]
bs)
      deps :: t -> [Int]
deps t
b = v -> Int
varId (v -> Int) -> [v] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set v -> [v]
forall a. Set a -> [a]
Set.toList (Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection Set v
vars (t -> Set v
freeVars t
b))
   in SCC (v, t) -> [(v, t)]
forall vertex. SCC vertex -> [vertex]
Graph.flattenSCC (SCC (v, t) -> [(v, t)]) -> [SCC (v, t)] -> [[(v, t)]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [((v, t), Int, [Int])] -> [SCC (v, t)]
forall key node. Ord key => [(node, key, [key])] -> [SCC node]
Graph.stronglyConnComp [((v, t), Int, [Int])]
graph