-- | The order to contract a network of factors in, such as a tensor network or the states of an
-- open hypergraph: which two to contract first, then which two of the factors and results, and so
-- on. Contracting two factors makes a factor on the nodes of both that something else still needs,
-- so its size depends on the order, and a bad order can make the intermediate results far larger
-- than the network itself.
--
-- 'planContraction' makes two plans and takes the one whose contractions are smaller together:
-- numpy's greedy order, and an elimination order of the nodes from the min-fill heuristic for tree
-- decompositions, as in Peterseim and Lopuhaä-Zwakenberg, /Fixed-Parameter Tractable Inference for
-- Discrete Probabilistic Programs, via String Diagram Algebraisation/ (LICS 2026).
module Proarrow.Tools.ContractionOrder (Plan (..), planContraction, minFillOrder) where

import Data.Containers.ListUtils (nubOrd)
import Data.List qualified as List
import Data.Map.Strict qualified as M
import Data.Ord (comparing)
import Data.Set (Set)
import Data.Set qualified as Set
import Prelude

-- | A plan to contract factors: one of them, by its place in the list, or the contraction of the
-- results of two plans.
data Plan = Factor Int | Contract Plan Plan
  deriving (Plan -> Plan -> Bool
(Plan -> Plan -> Bool) -> (Plan -> Plan -> Bool) -> Eq Plan
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Plan -> Plan -> Bool
== :: Plan -> Plan -> Bool
$c/= :: Plan -> Plan -> Bool
/= :: Plan -> Plan -> Bool
Eq, Int -> Plan -> ShowS
[Plan] -> ShowS
Plan -> String
(Int -> Plan -> ShowS)
-> (Plan -> String) -> ([Plan] -> ShowS) -> Show Plan
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Plan -> ShowS
showsPrec :: Int -> Plan -> ShowS
$cshow :: Plan -> String
show :: Plan -> String
$cshowList :: [Plan] -> ShowS
showList :: [Plan] -> ShowS
Show)

-- | A part of a plan: its plan, how many of its factors touch each node, the nodes something
-- outside it needs, and the size of its contractions together.
data Part n = Part {forall n. Part n -> Plan
plan :: Plan, forall n. Part n -> Map n Int
counts :: M.Map n Int, forall n. Part n -> Set n
open :: Set n, forall n. Part n -> Double
cost :: ~Double}

-- | Plans to contract the given factors, each given by its nodes, until no two of the results share
-- a node that may be contracted: a plan for each result. A result keeps the nodes that the given
-- set needs or that factors outside it touch, and its size is the size of those nodes. Of numpy's
-- greedy order, which contracts the pair that shares such a node and whose result grows least, its
-- size less those of the two parts, and of the min-fill elimination order of those nodes, which
-- contracts the parts on each node in turn, the pair that grows least first, the plan whose
-- contractions are smaller together is taken.
--
-- >>> planContraction (fromIntegral . length) (const True) mempty [Set.fromList "ab", Set.fromList "bc", Set.fromList "cd"]
-- [Contract (Contract (Factor 0) (Factor 1)) (Factor 2)]
planContraction :: (Ord n) => (Set n -> Double) -> (n -> Bool) -> Set n -> [Set n] -> [Plan]
planContraction :: forall n.
Ord n =>
(Set n -> Double) -> (n -> Bool) -> Set n -> [Set n] -> [Plan]
planContraction Set n -> Double
size n -> Bool
contractible Set n
needed [Set n]
factors
  -- with at most two factors on contractible nodes there is only one plan
  | [Int
i, Int
j] <- [Int]
candidates =
      if (n -> Bool) -> Set n -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any n -> Bool
contractible (Set n -> Set n -> Set n
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection ([Set n]
factors [Set n] -> Int -> Set n
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) ([Set n]
factors [Set n] -> Int -> Set n
forall a. HasCallStack => [a] -> Int -> a
!! Int
j))
        then Plan -> Plan -> Plan
Contract (Int -> Plan
Factor Int
i) (Int -> Plan
Factor Int
j) Plan -> [Plan] -> [Plan]
forall a. a -> [a] -> [a]
: Int -> Int -> [Plan]
others Int
i Int
j
        else [Plan]
alone
  | [Int
_] <- [Int]
candidates = [Plan]
alone
  | [] <- [Int]
candidates = [Plan]
alone
  | (Part n -> Bool) -> [Part n] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (Plan -> Bool
isFactor (Plan -> Bool) -> (Part n -> Plan) -> Part n -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Part n -> Plan
forall n. Part n -> Plan
plan) [Part n]
greedy = Part n -> Plan
forall n. Part n -> Plan
plan (Part n -> Plan) -> [Part n] -> [Plan]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [Part n]
greedy
  | Bool
otherwise = Part n -> Plan
forall n. Part n -> Plan
plan (Part n -> Plan) -> [Part n] -> [Plan]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> if [Part n] -> Double
forall {t :: Type -> Type} {n}.
(Foldable t, Functor t) =>
t (Part n) -> Double
total [Part n]
greedy Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= [Part n] -> Double
forall {t :: Type -> Type} {n}.
(Foldable t, Functor t) =>
t (Part n) -> Double
total [Part n]
eliminated then [Part n]
greedy else [Part n]
eliminated
  where
    candidates :: [Int]
candidates = [Int
i | (Int
i, Set n
f) <- [Int] -> [Set n] -> [(Int, Set n)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Set n]
factors, (n -> Bool) -> Set n -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any n -> Bool
contractible Set n
f]
    alone :: [Plan]
alone = Int -> Plan
Factor (Int -> Plan) -> [Int] -> [Plan]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. [Set n] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Set n]
factors Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    others :: Int -> Int -> [Plan]
others Int
i Int
j = [Int -> Plan
Factor Int
k | Int
k <- [Int
0 .. [Set n] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Set n]
factors Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
i, Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
j]
    leaves :: [Part n]
leaves = [Plan -> Map n Int -> Double -> Part n
part (Int -> Plan
Factor Int
i) ((n -> Int) -> Set n -> Map n Int
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (Int -> n -> Int
forall a b. a -> b -> a
const Int
1) Set n
f) Double
0 | (Int
i, Set n
f) <- [Int] -> [Set n] -> [(Int, Set n)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Set n]
factors]
    greedy :: [Part n]
greedy = [Part n] -> [Part n]
contract [Part n]
leaves
    eliminated :: [Part n]
eliminated = [n] -> [Part n] -> [Part n]
along (([n] -> Double) -> [[n]] -> [n] -> [n]
forall a w. (Ord a, Ord w) => ([a] -> w) -> [[a]] -> [a] -> [a]
minFillOrder (Set n -> Double
size (Set n -> Double) -> ([n] -> Set n) -> [n] -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [n] -> Set n
forall a. Ord a => [a] -> Set a
Set.fromList) (Set n -> [n]
forall a. Set a -> [a]
Set.toList (Set n -> [n]) -> (Part n -> Set n) -> Part n -> [n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Part n -> Set n
forall n. Part n -> Set n
open (Part n -> [n]) -> [Part n] -> [[n]]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [Part n]
leaves) [n]
nodes) [Part n]
leaves
    total :: t (Part n) -> Double
total t (Part n)
ps = t Double -> Double
forall a. Num a => t a -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum (Part n -> Double
forall n. Part n -> Double
cost (Part n -> Double) -> t (Part n) -> t Double
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> t (Part n)
ps)
    nodes :: [n]
nodes = [n] -> [n]
forall a. Ord a => [a] -> [a]
nubOrd [n
n | Set n
f <- [Set n]
factors, n
n <- Set n -> [n]
forall a. Set a -> [a]
Set.toList Set n
f, n -> Bool
contractible n
n]
    totals :: Map n Int
totals = (Int -> Int -> Int) -> [(n, Int)] -> Map n Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(n
n, Int
1) | Set n
f <- [Set n]
factors, n
n <- Set n -> [n]
forall a. Set a -> [a]
Set.toList Set n
f]
    part :: Plan -> Map n Int -> Double -> Part n
part Plan
p Map n Int
cs Double
c = Plan -> Map n Int -> Set n -> Double -> Part n
forall n. Plan -> Map n Int -> Set n -> Double -> Part n
Part Plan
p Map n Int
cs (Map n Int -> Set n
forall k a. Map k a -> Set k
M.keysSet ((n -> Int -> Bool) -> Map n Int -> Map n Int
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\n
n Int
k -> Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Map n Int
totals Map n Int -> n -> Int
forall k a. Ord k => Map k a -> k -> a
M.! n
n Bool -> Bool -> Bool
|| n
n n -> Set n -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set n
needed) Map n Int
cs)) Double
c
    merge :: Part n -> Part n -> Part n
merge Part n
a Part n
b =
      Plan -> Map n Int -> Double -> Part n
part
        (Plan -> Plan -> Plan
Contract (Part n -> Plan
forall n. Part n -> Plan
plan Part n
a) (Part n -> Plan
forall n. Part n -> Plan
plan Part n
b))
        ((Int -> Int -> Int) -> Map n Int -> Map n Int -> Map n Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Part n -> Map n Int
forall n. Part n -> Map n Int
counts Part n
a) (Part n -> Map n Int
forall n. Part n -> Map n Int
counts Part n
b))
        (Part n -> Double
forall n. Part n -> Double
cost Part n
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Part n -> Double
forall n. Part n -> Double
cost Part n
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Set n -> Double
size (Set n -> Set n -> Set n
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Part n -> Set n
forall n. Part n -> Set n
open Part n
a) (Part n -> Set n
forall n. Part n -> Set n
open Part n
b)))
    isFactor :: Plan -> Bool
isFactor Factor{} = Bool
True
    isFactor Contract{} = Bool
False
    shares :: Part n -> Part n -> Bool
shares Part n
a Part n
b = (n -> Bool) -> Set n -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any n -> Bool
contractible (Set n -> Set n -> Set n
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection (Part n -> Set n
forall n. Part n -> Set n
open Part n
a) (Part n -> Set n
forall n. Part n -> Set n
open Part n
b))
    pairsOf :: [b] -> [(b, b)]
pairsOf [b]
ps = [(b
a, b
b) | b
a : [b]
rest <- [b] -> [[b]]
forall a. [a] -> [[a]]
List.tails [b]
ps, b
b <- [b]
rest]
    -- of the given pairs, the one whose contraction grows least, with that contraction
    best :: [(Part n, Part n)] -> Maybe (Part n, Part n, Part n)
best [(Part n, Part n)]
pairs = case [(Set n -> Double
size (Part n -> Set n
forall n. Part n -> Set n
open Part n
ab) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Set n -> Double
size (Part n -> Set n
forall n. Part n -> Set n
open Part n
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Set n -> Double
size (Part n -> Set n
forall n. Part n -> Set n
open Part n
b), (Part n
a, Part n
b, Part n
ab)) | (Part n
a, Part n
b) <- [(Part n, Part n)]
pairs, let ab :: Part n
ab = Part n -> Part n -> Part n
merge Part n
a Part n
b] of
      [] -> Maybe (Part n, Part n, Part n)
forall a. Maybe a
Nothing
      [(Double, (Part n, Part n, Part n))]
cs -> (Part n, Part n, Part n) -> Maybe (Part n, Part n, Part n)
forall a. a -> Maybe a
Just ((Double, (Part n, Part n, Part n)) -> (Part n, Part n, Part n)
forall a b. (a, b) -> b
snd (((Double, (Part n, Part n, Part n))
 -> (Double, (Part n, Part n, Part n)) -> Ordering)
-> [(Double, (Part n, Part n, Part n))]
-> (Double, (Part n, Part n, Part n))
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
List.minimumBy (((Double, (Part n, Part n, Part n)) -> Double)
-> (Double, (Part n, Part n, Part n))
-> (Double, (Part n, Part n, Part n))
-> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Double, (Part n, Part n, Part n)) -> Double
forall a b. (a, b) -> a
fst) [(Double, (Part n, Part n, Part n))]
cs))
    replaced :: Part n -> Part n -> Part n -> [Part n] -> [Part n]
replaced Part n
a Part n
b Part n
ab [Part n]
ps = [if Part n -> Plan
forall n. Part n -> Plan
plan Part n
q Plan -> Plan -> Bool
forall a. Eq a => a -> a -> Bool
== Part n -> Plan
forall n. Part n -> Plan
plan Part n
a then Part n
ab else Part n
q | Part n
q <- [Part n]
ps, Part n -> Plan
forall n. Part n -> Plan
plan Part n
q Plan -> Plan -> Bool
forall a. Eq a => a -> a -> Bool
/= Part n -> Plan
forall n. Part n -> Plan
plan Part n
b]
    contract :: [Part n] -> [Part n]
contract [Part n]
ps = case [(Part n, Part n)] -> Maybe (Part n, Part n, Part n)
best [(Part n
a, Part n
b) | (Part n
a, Part n
b) <- [Part n] -> [(Part n, Part n)]
forall {b}. [b] -> [(b, b)]
pairsOf [Part n]
ps, Part n -> Part n -> Bool
shares Part n
a Part n
b] of
      Just (Part n
a, Part n
b, Part n
ab) -> [Part n] -> [Part n]
contract (Part n -> Part n -> Part n -> [Part n] -> [Part n]
forall {n} {n} {n}.
Part n -> Part n -> Part n -> [Part n] -> [Part n]
replaced Part n
a Part n
b Part n
ab [Part n]
ps)
      Maybe (Part n, Part n, Part n)
Nothing -> [Part n]
ps
    along :: [n] -> [Part n] -> [Part n]
along [] [Part n]
ps = [Part n]
ps
    along (n
n : [n]
ns) [Part n]
ps = case [(Part n, Part n)] -> Maybe (Part n, Part n, Part n)
best ([Part n] -> [(Part n, Part n)]
forall {b}. [b] -> [(b, b)]
pairsOf [Part n
p | Part n
p <- [Part n]
ps, n
n n -> Set n -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Part n -> Set n
forall n. Part n -> Set n
open Part n
p]) of
      Just (Part n
a, Part n
b, Part n
ab) -> [n] -> [Part n] -> [Part n]
along (n
n n -> [n] -> [n]
forall a. a -> [a] -> [a]
: [n]
ns) (Part n -> Part n -> Part n -> [Part n] -> [Part n]
forall {n} {n} {n}.
Part n -> Part n -> Part n -> [Part n] -> [Part n]
replaced Part n
a Part n
b Part n
ab [Part n]
ps)
      Maybe (Part n, Part n, Part n)
Nothing -> [n] -> [Part n] -> [Part n]
along [n]
ns [Part n]
ps

-- | An order to eliminate the given vertices of a graph in, by the min-fill heuristic for tree
-- decompositions. The graph has an edge between any two vertices of a common clique. Eliminating a
-- vertex joins its neighbours, so the next vertex is the one whose neighbours lack the fewest edges
-- among themselves, and of those the one whose neighbours weigh least together. A star is
-- eliminated from its leaves:
--
-- >>> minFillOrder length [[0, 1], [0, 2], [0, 3 :: Int]] [0, 1, 2, 3]
-- [1,2,0,3]
minFillOrder :: (Ord a, Ord w) => ([a] -> w) -> [[a]] -> [a] -> [a]
minFillOrder :: forall a w. (Ord a, Ord w) => ([a] -> w) -> [[a]] -> [a] -> [a]
minFillOrder [a] -> w
weight [[a]]
cliques = Map a (Set a) -> [a] -> [a]
go Map a (Set a)
graph
  where
    graph :: Map a (Set a)
graph = (Set a -> Set a -> Set a) -> [(a, Set a)] -> Map a (Set a)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union [(a
v, [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
v) [a]
c)) | [a]
c <- [[a]]
cliques, a
v <- [a]
c]
    neighbours :: Map k (Set a) -> k -> Set a
neighbours Map k (Set a)
g k
v = Set a -> k -> Map k (Set a) -> Set a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set a
forall a. Set a
Set.empty k
v Map k (Set a)
g
    go :: Map a (Set a) -> [a] -> [a]
go Map a (Set a)
_ [] = []
    go Map a (Set a)
g [a]
vs =
      let v :: a
v = (a -> a -> Ordering) -> [a] -> a
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
List.minimumBy ((a -> (Int, w)) -> a -> a -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Map a (Set a) -> a -> (Int, w)
fill Map a (Set a)
g)) [a]
vs
          nbs :: Set a
nbs = Map a (Set a) -> a -> Set a
forall {k} {a}. Ord k => Map k (Set a) -> k -> Set a
neighbours Map a (Set a)
g a
v
          -- only the neighbours of v change: they lose v and become adjacent to each other
          joined :: Map a (Set a)
joined = (a -> Map a (Set a) -> Map a (Set a))
-> Map a (Set a) -> [a] -> Map a (Set a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
u -> (Set a -> Set a) -> a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
u (Set a -> Set a) -> (Set a -> Set a) -> Set a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
nbs (Set a -> Set a) -> (Set a -> Set a) -> Set a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
v) a
u) (a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
v Map a (Set a)
g) (Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
nbs)
      in a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Map a (Set a) -> [a] -> [a]
go Map a (Set a)
joined (a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
List.delete a
v [a]
vs)
    fill :: Map a (Set a) -> a -> (Int, w)
fill Map a (Set a)
g a
v =
      let nbs :: [a]
nbs = Set a -> [a]
forall a. Set a -> [a]
Set.toList (Map a (Set a) -> a -> Set a
forall {k} {a}. Ord k => Map k (Set a) -> k -> Set a
neighbours Map a (Set a)
g a
v)
      in ([()] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [() | (a
a : [a]
rest) <- [a] -> [[a]]
forall a. [a] -> [[a]]
List.tails [a]
nbs, let na :: Set a
na = Map a (Set a) -> a -> Set a
forall {k} {a}. Ord k => Map k (Set a) -> k -> Set a
neighbours Map a (Set a)
g a
a, a
b <- [a]
rest, Bool -> Bool
not (a
b a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a
na)], [a] -> w
weight [a]
nbs)