srtree 2.0.1.8 → 3.0.0.0

raw patch · 47 files changed

+7889/−5811 lines, 47 filesdep +asyncdep +criteriondep +deepseqdep −dlistdep −list-shuffledep −massivdep ~HUnitdep ~addep ~binarynew-component:exe:benchnew-component:exe:bench-eqsatnew-component:exe:srtree-reportPVP ok

version bump matches the API change (PVP)

Dependencies added: async, criterion, deepseq, directory, parallel, primitive, time

Dependencies removed: dlist, list-shuffle, massiv, scheduler, transformers, unliftio, unliftio-core

Dependency ranges changed: HUnit, ad, binary, exceptions, hashable, ieee754, lens, optparse-applicative, statistics

API changes (from Hackage documentation)

Files

• ChangeLog.md +15/−0
• LICENSE +3/−4
• apps/Bench/Main.hs +127/−0
• apps/BenchEqSat/Main.hs +259/−0
• apps/Report/Main.hs +247/−0
• apps/srsimplify/Main.hs +0/−103
• apps/srtools/Args.hs +0/−184
• apps/srtools/IO.hs +0/−230
• apps/srtools/Main.hs +0/−34
• apps/srtools/Report.hs +0/−280
• apps/tinygp/GP.hs +0/−252
• apps/tinygp/Initialization.hs +0/−50
• apps/tinygp/Main.hs +0/−116
• apps/tinygp/Util.hs +0/−63
• src/Algorithm/EqSat.hs +290/−96
• src/Algorithm/EqSat/Build.hs +432/−329
• src/Algorithm/EqSat/DB.hs +421/−99
• src/Algorithm/EqSat/Egraph.hs +741/−174
• src/Algorithm/EqSat/Info.hs +86/−72
• src/Algorithm/EqSat/Queries.hs +100/−108
• src/Algorithm/EqSat/SearchSR.hs +115/−61
• src/Algorithm/EqSat/SearchSRCache.hs +0/−244
• src/Algorithm/EqSat/Simplify.hs +117/−113
• src/Algorithm/EqSat/Store.hs +243/−0
• src/Algorithm/Massiv/Utils.hs +0/−278
• src/Algorithm/SRTree/AD.hs +14/−541
• src/Algorithm/SRTree/AD/CompiledAD.hs +39/−0
• src/Algorithm/SRTree/AD/Unboxed.hs +965/−0
• src/Algorithm/SRTree/Compile.hs +106/−0
• src/Algorithm/SRTree/ConfidenceIntervals.hs +248/−274
• src/Algorithm/SRTree/Likelihoods.hs +216/−522
• src/Algorithm/SRTree/ModelSelection.hs +127/−119
• src/Algorithm/SRTree/NonlinearOpt.hs +98/−976
• src/Algorithm/SRTree/Opt.hs +0/−135
• src/Algorithm/SRTree/Utils.hs +320/−0
• src/Data/SRTree/Datasets.hs +188/−53
• src/Data/SRTree/Derivative.hs +13/−0
• src/Data/SRTree/Eval.hs +228/−50
• src/Data/SRTree/Internal.hs +10/−6
• src/Data/SRTree/Random.hs +3/−3
• src/Numeric/Optimization/NLOPT.hs +976/−0
• src/Text/ParseSR.hs +7/−7
• src/Text/ParseSR/IO.hs +2/−2
• srtree.cabal +221/−231
• test/EqSatTests.hs +630/−0
• test/Spec.hs +113/−2
• test/StoreTests.hs +169/−0

@@ -1,5 +1,20 @@ # Changelog for srtree +## 3.0.0.0++- **BREAKING**: Removed the Accelerate AD backend (`Algorithm.SRTree.AD.Accelerate`).+  The `ADBackEnd` type now only has `SingleThread` and `MultiThread` constructors.+  This removes the `accelerate` and `accelerate-llvm-native` dependencies.+- Out-of-core equality saturation with paged e-graph store (SQLite/PostgreSQL)+- Frontier re-saturation: mark changed classes and re-saturate only the frontier+- Streaming matcher for n-ary and cached genericJoin paths (O(1) memory on paged graphs)+- Cycle-safe and size-budgeted `getBestExpr` extraction+- Bounded cost/best fixpoints so recalc terminates on cyclic graphs+- Bounded node-to-class and canonical maps on paged graphs (LRU caches)+- Fast ByteString double parser for dataset loading+- Thread `Loss` (not `Distribution`) through fitness functions; add `readLoss`+- Multiset e-graph improvements+ ## 2.0.1.7   - Added log10 MSE fitness function 

@@ -1,6 +1,5 @@-Copyright Author name here (c) 2021+Copyright (c) 2026, folivetti -All rights reserved.  Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:@@ -13,7 +12,7 @@       disclaimer in the documentation and/or other materials provided       with the distribution. -    * Neither the name of Author name here nor the names of other+    * Neither the name of the copyright holder nor the names of its       contributors may be used to endorse or promote products derived       from this software without specific prior written permission. @@ -21,7 +20,7 @@ "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY

@@ -0,0 +1,127 @@+{-# LANGUAGE BangPatterns #-}++import Criterion.Main+import Control.DeepSeq (force, NFData)+import Control.Exception (evaluate)+import qualified Data.Vector.Unboxed as V+import qualified Data.Vector as VB+import qualified Data.Vector.Generic as G+import qualified Data.Vector.Storable as VS++import Data.SRTree+import Data.SRTree.Print+import Data.SRTree.Datasets+import Data.SRTree.Eval+import Data.SRTree.Random+import System.Random+import Control.Monad.State.Strict+import Algorithm.SRTree.NonlinearOpt+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.AD++-- Assuming these are exported by your project modules:+-- import SRTree+-- import Compiler+-- import DatasetLoader++-- Mock signatures based on your provided functions+-- randomTree :: Int -> Int -> Int -> IO Term -> IO NonTerm -> Bool -> IO Tree+-- loadDataset :: FilePath -> Bool -> IO [V.Vector Double]+-- evalTree :: Tree -> [V.Vector Double] -> V.Vector Double+-- compile :: [V.Vector Double] -> Tree -> (Theta -> V.Vector Double)++genTerm    = do coin <- tossBiased 0.4+                if coin then randomFrom [Fix $ Var ix | ix <- [0..8]] else randomFrom [Fix $ Param ix | ix <- [0..9]]+genNonTerm = randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Uni LogAbs (), Uni SqrtAbs ()]++genMultipleTrees 0 = pure []+genMultipleTrees n = do+    t <- randomTree 5 10 150 genTerm genNonTerm False+    ts <- genMultipleTrees (n-1)+    pure (t:ts)++getF (_, x, _) = x+{-# INLINE getF #-}+getT (t, _, _) = t+{-# INLINE getT #-}++main :: IO ()+main = do+    -- 1. Initialization: Load the dataset+    putStrLn "Loading dataset..."+    ((dataset, y, _, _), _, _, _) <- loadDataset "data.tsv" True++    -- 2. Initialization: Generate the random expression tree+    putStrLn "Generating random tree..."+    -- Replace 'genTerm' and 'genNonTerm' with your actual generators+    --g <- getStdGen+    let g = mkStdGen 42+    -- tree <- evalStateT (randomTree 7 10 150 genTerm genNonTerm True) g+    trees' <- evalStateT (genMultipleTrees 5) g+    -- let trees' = [Fix (Uni LogAbs (Fix (Bin PowerAbs (param 0) (param 1 * var 0))))] :: [Fix SRTree]++    -- IMPORTANT: Force deep evaluation of the tree and dataset.+    -- If we do not do this, GHC's lazy evaluation will cause the benchmark+    -- to measure the time it takes to parse the CSV and build the tree in memory!+    -- _ <- evaluate (force tree)+    _ <- evaluate (force dataset)+++    -- 3. Initialization: Pre-compile the tree+    -- We evaluate this strictly (!) so the one-time compilation cost+    -- is not included in the runtime benchmark.+    putStrLn "Compiling tree..."+    let !compiledFn = [compile dataset tree | tree <- trees]+        evalTree x th t = compile x t th+        -- Mock theta (parameter vector) to pass into the closures+        !theta = V.fromList [1.0, 0.5, 0.2, 0.3, 0.1, 0.5, 0.9, 0.3, 0.2, 0.4]+        !theta1 = V.fromList [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]+        trees = map relabelParamsOrder $ filter (\t -> let v = V.sum (evalTree dataset theta t) in not (isInfinite v || isNaN v)) trees'+        naiveEval = evalTree dataset theta+        dataset' = map G.convert dataset+        y' = G.convert y+        theta1' = G.convert theta1++    _ <- evaluate (force theta)+    _ <- evaluate (force theta1)+    print $ sum $ map (\t ->  V.sum $ naiveEval t) trees+    print $ sum $ map (\t ->  V.sum $ t theta) compiledFn+    print $ sum $ map (\t -> getF $ minimizeNLL MultiThread MSE Nothing 0 dataset y t theta1) trees+    --print $ sum $ map (\t -> getF $ minimizeNLLCompiled MSE Nothing 0 dataset y t theta1) trees++    --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraph MSE dataset' y' Nothing t theta1') trees+    --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraphO MSE dataset' y' Nothing t theta1') trees+    --print $ sum $ map (\t -> VS.sum . snd $ compileGrad dataset' y' Nothing t 100 theta1') trees+    --print $ sum $ map (\ct -> V.sum $ ct theta) compiledFn+    --print $ sum $ map (\ct -> V.sum $ executeVM ct rowDataset theta) bytecodes+    -- print $ V.sum $ evalTree dataset theta tree+    -- print $ V.sum $ compiledFn theta++    putStrLn "Running benchmarks..."++    -- 4. The Benchmarks+    defaultMain [+          bgroup "Tree Evaluation (Fixed Dataset)" [++           -- The slow version: dynamically traversing the AST at runtime+           bench "evalTree (Naive AST Traversal)" $+                nf (\ts -> sum [V.sum $ evalTree dataset theta1 t | t <- ts]) trees,+++            -- The fast version: executing the pre-compiled, stream-fused closure+            bench "compile (Compiled Closure)" $+                nf (\t -> sum [V.sum (ct t) | ct <- compiledFn]) theta1,++            -- The fast version: executing the pre-compiled, stream-fused closure+            bench "minimizeNLLCompiled (Compiled Closure)" $+                nf (\ts -> sum [V.sum . getT $ minimizeNLL MultiThread MSE Nothing 100 dataset' y' t theta1' | t <- ts]) trees++            --bench "minimizeNLLO (Naive optimized AST Traversal)" $+            --    nf (\ts -> sum [V.sum . getT $ minimizeNLLO MSE Nothing 100 dataset y t theta1 | t <- ts]) trees++            -- The slow version: dynamically traversing the AST at runtime+            --bench "minimizeNLL (Naive AST Traversal)" $+            --    nf (\ts -> sum [V.sum . getT $ minimizeNLL (NLL MSE) Nothing 100 dataset y t theta1 | t <- ts]) trees++        ]+      ]

@@ -0,0 +1,259 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedStrings #-}++import Criterion.Main+import qualified Data.Vector.Unboxed as VU+import qualified Data.IntMap as IntMap+import qualified Data.HashMap.Strict as HashMap+import qualified Data.HashSet as Set++import Data.SRTree+import Algorithm.EqSat+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.DB+import Algorithm.EqSat.Info+import Algorithm.EqSat.Queries+import Control.Monad.State.Strict+import Control.Monad (replicateM, zipWithM_)+import Control.Monad.Identity++myCost :: SRTree Int -> Int+myCost (Var _)     = 1+myCost (Const _)   = 1+myCost (Param _)   = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t)   = 3 + t++evalEG :: EGraphST Identity a -> (a, EGraph)+evalEG m = runIdentity $ runStateT m emptyGraph++runInEG :: EGraph -> EGraphST Identity a -> (a, EGraph)+runInEG eg m = runIdentity $ runStateT m eg++-- Expression generators for benchmarking+chainAdd :: Int -> Fix SRTree+chainAdd 0 = var 0+chainAdd n = chainAdd (n-1) + var n++deepBinTree :: Int -> Fix SRTree+deepBinTree 0 = var 0+deepBinTree n = deepBinTree (n-1) + constv (fromIntegral n)++complexTree :: Int -> Fix SRTree+complexTree n = go n+  where+    go 0 = var 0+    go i = (var i + constv (fromIntegral i)) * (go (i-1) + constv (fromIntegral i))++simplifyRules :: [Rule]+simplifyRules =+  [ "a" + 0 :=> "a"+  , "a" * 1 :=> "a"+  , "a" + "a" :=> 2 * "a"+  , "a" * 0 :=> 0+  , 0 + "a" :=> "a"+  , 1 * "a" :=> "a"+  ]++-- More rules including commutativity (triggers more merges)+moreRules :: [Rule]+moreRules =+  [ "a" + 0 :=> "a"+  , "a" * 1 :=> "a"+  , "a" + "a" :=> 2 * "a"+  , "a" * 0 :=> 0+  , 0 + "a" :=> "a"+  , 1 * "a" :=> "a"+  , "a" + "b" :=> "b" + "a"+  , "a" * "b" :=> "b" * "a"+  ]++addZero :: Fix SRTree -> Fix SRTree -> Fix SRTree+addZero l r = Fix (Bin Add l r)++main :: IO ()+main = do+  putStrLn "Generating benchmark expressions..."+  let smallExpr  = chainAdd 5+      mediumExpr = chainAdd 20+      largeExpr  = chainAdd 100+      complex    = complexTree 8++  putStrLn "Running benchmarks..."+  defaultMain [+    bgroup "E-graph Construction" [+      bench "fromTree (5-leaf chain)" $+        whnf (\e -> evalEG $ fromTree myCost e) smallExpr,+      bench "fromTree (20-leaf chain)" $+        whnf (\e -> evalEG $ fromTree myCost e) mediumExpr,+      bench "fromTree (100-leaf chain)" $+        whnf (\e -> evalEG $ fromTree myCost e) largeExpr,+      bench "fromTree (complex-ternary tree)" $+        whnf (\e -> evalEG $ fromTree myCost e) complex+    ],++    bgroup "E-graph Add" [+      bench "add single e-node (Var)" $+        whnf (\eg -> runInEG eg $ add myCost (EVar 999)) (snd $ evalEG $ fromTree myCost smallExpr),+      bench "add single e-node (Const)" $+        whnf (\eg -> runInEG eg $ add myCost (EConst 42.0)) (snd $ evalEG $ fromTree myCost smallExpr),+      bench "add single e-node (Bin Add)" $+        whnf (\eg -> runInEG eg $ add myCost (ENAry EAdd (imFromList [0, 1]))) (snd $ evalEG $ fromTree myCost mediumExpr)+    ],++    bgroup "Merge" [+      bench "merge two distinct eclasses (size 1)" $+        whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 1),+      bench "merge two distinct eclasses (size 3)" $+        whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 3)+    ],++    bgroup "Pattern Matching" [+      bench "match simple pattern (a+0)" $+        whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMatchableEG),+      bench "match commutative pattern (a+b)" $+        whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMatchableEG),+      bench "match triple pattern (a+b+c)" $+        whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" + "c" :: Pattern)) (makeMatchableEG)+    ],++    bgroup "Match After Merge" [+      bench "match (a+0) after merge (stale trie keys)" $+        whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMergedEG),+      bench "match (a+b) after merge (stale trie keys)" $+        whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMergedEG)+    ],++    bgroup "Rebuild" [+      bench "rebuild after 5 adds" $+        whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 5),+      bench "rebuild after 20 adds" $+        whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 20),+      bench "rebuild after 100 adds" $+        whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 100)+    ],++    bgroup "Cost Propagation" [+      bench "recalculateBest (10 eclasses)" $+        whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 10),+      bench "recalculateBest (100 eclasses)" $+        whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 100)+    ],++    bgroup "DB Operations" [+      bench "addToDB single enode" $+        whnf (\(en,eid,eg) -> runInEG eg $ addToDB en eid) (makeDBEntry),+      bench "addToDB 10 enodes" $+        whnf (\(ens,eg) -> runInEG eg $ mapM_ (uncurry addToDB) ens) (makeDBEntries 10)+    ],++    bgroup "Equality Saturation" [+      bench "eqSat small expr (5 rules)" $+        whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, simplifyRules),+      bench "eqSat medium expr (5 rules)" $+        whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (mediumExpr, simplifyRules),+      bench "eqSat small expr (8 rules, commutative)" $+        whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, moreRules),+      bench "eqSat large expr (5 rules)" $+        whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (largeExpr, simplifyRules)+    ],++    bgroup "Extraction" [+      bench "getBestExpr (5-leaf)" $+        whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 5),+      bench "getBestExpr (20-leaf)" $+        whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 20),+      bench "getBestExpr (100-leaf)" $+        whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 100)+    ],++    bgroup "Fitness Operations" [+      bench "insertFitness single" $+        whnf (\(eid,eg) -> runInEG eg $ insertFitness eid 0.5 []) (makeExtractable 1),+      bench "insertFitness 10 eclasses" $+        whnf (\(eids,eg) -> runInEG eg $ mapM_ (\eid -> insertFitness eid 0.5 []) eids) (makeNEclasses 10),+      bench "getTopFitEClassWithSize" $+        whnf (\(eids,eg) -> runInEG eg $ getTopFitEClassWithSize 1 3) (makeFitnessEG)+    ]+    ]+  where+    addZeroTree = addZero (var 0) (constv 0.0)++    makeMergePair :: Int -> (EClassId, EClassId, EGraph)+    makeMergePair n =+      let tree = deepBinTree n+          (eid1, eg1) = evalEG $ fromTree myCost tree+          (eid2, eg2) = runInEG eg1 $ fromTree myCost tree+      in (eid1, eid2, eg2)++    makeMatchableEG :: (EGraph, EClassId)+    makeMatchableEG =+      let tree = complexTree 4+          (eid, eg) = evalEG $ do+            eid' <- fromTree myCost tree+            _ <- fromTree myCost (var 0 + constv 1.0)+            _ <- fromTree myCost (var 1 * constv 2.0)+            _ <- fromTree myCost (var 0 + constv 0.0)+            _ <- fromTree myCost (var 1 * constv 1.0)+            rebuild myCost+            pure eid'+      in (eg, eid)++    -- E-graph with merges applied, creating stale trie keys+    makeMergedEG :: (EGraph, EClassId)+    makeMergedEG =+      let (_, eg) = evalEG $ do+            eid1 <- fromTree myCost (var 0)+            eid2 <- fromTree myCost (constv 0.0)+            eid3 <- fromTree myCost (var 0 + constv 1.0)+            _ <- fromTree myCost (var 1)+            rebuild myCost+            -- merge to create stale trie entries+            merge myCost eid1 eid2+            merge myCost eid2 eid3+            rebuild myCost+            pure eid1+      in (eg, 0)++    makeDirtyEG :: Int -> (EGraph, EClassId)+    makeDirtyEG n =+      let tree = deepBinTree n+          (eid, eg) = evalEG $ do+            eid' <- fromTree myCost tree+            _ <- fromTree myCost (tree + var 999)+            rebuild myCost+            _ <- fromTree myCost (tree * var 998)+            pure eid'+      in (eg, eid)++    makeExtractable :: Int -> (EClassId, EGraph)+    makeExtractable n =+      let tree = deepBinTree n+      in evalEG $ fromTree myCost tree++    makeNEclasses :: Int -> ([EClassId], EGraph)+    makeNEclasses n =+      evalEG $ replicateM n (fromTree myCost (constv (fromIntegral n)))++    makeFitnessEG :: ([EClassId], EGraph)+    makeFitnessEG = evalEG $ do+      eids <- mapM (fromTree myCost . constv . fromIntegral) [1..10]+      zipWithM_ (\eid i -> insertFitness eid (fromIntegral i) []) eids [1..]+      pure eids++    makeDBEntry :: (ENode, EClassId, EGraph)+    makeDBEntry =+      let (eid, eg) = evalEG $ do+            eid <- fromTree myCost (var 999)+            rebuild myCost+            pure eid+      in (EVar 777, eid, eg)++    makeDBEntries :: Int -> ([(ENode, EClassId)], EGraph)+    makeDBEntries n =+      let (eids, eg) = evalEG $ do+            eids <- mapM (fromTree myCost . var) [999..(999 + n - 1)]+            rebuild myCost+            pure eids+      in (zip (map EVar [1000..]) eids, eg)

@@ -0,0 +1,247 @@+module Main (main) where++import Options.Applicative+import qualified Data.ByteString.Char8 as B+import qualified Data.Vector.Unboxed as U+import Data.SRTree+import Data.SRTree.Eval (Target, Columns, compileLoss)+import Data.SRTree.Datasets (loadTrainingOnly)+import Data.SRTree.Print (showExpr)+import Text.ParseSR (parseSR, SRAlgs(..))+import Algorithm.SRTree.Compile (compileTree, EvalTree(..), logParameters, logParametersLatt)+import Algorithm.SRTree.Likelihoods (Distribution(..), Loss(..), buildLoss, fisherNLL, hessianNLL)+import Algorithm.SRTree.ConfidenceIntervals+    ( getStatsFromModel, paramCI, CIType(..), CI(..), BasicStats(..)+    , ProfileT(..), PType(..), getAllProfiles, getCol+    )+import Algorithm.SRTree.ModelSelection (ModelEval(..), logFunctional, logFunctionalFreq)+import Statistics.Distribution (ContDistr(quantile))+import Statistics.Distribution.FDistribution (fDistribution)+import Control.Exception (try, SomeException)+import Data.List.Split (splitOn)+import Text.Printf (printf)+import Control.Monad (forM_, when)++----------------------------------------------------------------------+-- CLI argument types+----------------------------------------------------------------------+data CIMethod = LaplaceCI | ProfileCI deriving (Show)++data ProfileTypeArg = BatesArg | ODEArg | ConstrainedArg deriving (Read)+instance Show ProfileTypeArg where+  show BatesArg       = "Bates"+  show ODEArg         = "ODE"+  show ConstrainedArg = "Constrained"++data ReportArgs = ReportArgs+  { raExprs    :: !FilePath+  , raFormat   :: !SRAlgs+  , raData     :: !FilePath+  , raHeader   :: !Bool+  , raDist     :: !Distribution+  , raCriteria :: ![ModelEval]+  , raCI       :: !CIMethod+  , raAlpha    :: !Double+  , raCIType   :: !ProfileTypeArg+  , raDbg      :: !Bool+  }++----------------------------------------------------------------------+-- Argument parser+----------------------------------------------------------------------+argParser :: Parser ReportArgs+argParser = ReportArgs+  <$> strOption ( long "exprs" <> short 'e' <> help "File with expressions, one per line" <> metavar "FILE" )+  <*> option auto ( long "format" <> short 'f' <> help "Expression format: TIR, HL, OPERON, BINGO, GOMEA, PYSR, SBP, EPLEX" <> metavar "FMT" )+  <*> strOption ( long "data" <> short 'd' <> help "Dataset file (optionally with :start:end:target:features:y_err)" <> metavar "FILE" )+  <*> switch ( long "header" <> help "Dataset has a header row" )+  <*> option auto ( long "dist" <> value Gaussian <> help "Distribution: Gaussian, Bernoulli, Poisson, LeastSquares" <> metavar "DIST" <> showDefault )+  <*> option parseCriteria ( long "criteria" <> short 'c' <> value [RMSE, R2, AIC, BIC] <> help "Comma-separated criteria" <> metavar "CRITERIA" <> showDefault )+  <*> option parseCI ( long "ci" <> value LaplaceCI <> help "CI method: Laplace, Profile" <> metavar "METHOD" <> showDefault )+  <*> option auto ( long "alpha" <> value 0.05 <> help "Significance level" <> metavar "ALPHA" <> showDefault )+  <*> option parseProfileType ( long "ci-type" <> value BatesArg <> help "Profile CI type: Bates, ODE, Constrained" <> metavar "TYPE" <> showDefault )+  <*> switch ( long "dbg" <> help "Debug: dump profile tau/theta spline points" )++parseCriteria :: ReadM [ModelEval]+parseCriteria = eitherReader $ \s ->+  case traverse parseOne (splitOn "," s) of+    Right es -> Right es+    Left  e  -> Left e+  where+    parseOne "RMSE"     = Right RMSE+    parseOne "R2"       = Right R2+    parseOne "AIC"      = Right AIC+    parseOne "BIC"      = Right BIC+    parseOne "Evidence" = Right Evidence+    parseOne "FBF"      = Right FBF+    parseOne "MDL"      = Right MDL+    parseOne "MDLLatt"  = Right MDLLatt+    parseOne "MDLFreq"  = Right MDLFreq+    parseOne "NLL"      = Right (EvalLoss (NLL Gaussian))+    parseOne s          = Left ("unknown criterion: " ++ s)++parseCI :: ReadM CIMethod+parseCI = eitherReader $ \s -> case s of+  "Laplace" -> Right LaplaceCI+  "Profile" -> Right ProfileCI+  _         -> Left ("unknown CI method: " ++ s ++ " (use Laplace or Profile)")++parseProfileType :: ReadM ProfileTypeArg+parseProfileType = eitherReader $ \s -> case s of+  "Bates"       -> Right BatesArg+  "ODE"         -> Right ODEArg+  "Constrained" -> Right ConstrainedArg+  _             -> Left ("unknown profile type: " ++ s ++ " (use Bates, ODE, or Constrained)")++----------------------------------------------------------------------+-- Report data+----------------------------------------------------------------------+data ReportData = ReportData+  { rdTree      :: Fix SRTree+  , rdTheta     :: Target+  , rdStdErr    :: Target+  , rdCriteria  :: [(ModelEval, Double)]+  , rdCIs       :: [CI]+  }++----------------------------------------------------------------------+-- Main+----------------------------------------------------------------------+main :: IO ()+main = do+  args <- execParser (info (argParser <**> helper) fullDesc)+  (xss, ys, mYerr) <- loadTrainingOnly (raData args) (raHeader args)+  content <- B.readFile (raExprs args)+  let exprs = filter (not . B.null) $ B.lines content+  mapM_ (processOne args xss ys mYerr) (zip [(1 :: Int) ..] exprs)++----------------------------------------------------------------------+-- Process a single expression+----------------------------------------------------------------------+processOne :: ReportArgs -> Columns -> Target -> Maybe Target -> (Int, B.ByteString) -> IO ()+processOne args xss ys mYerr (idx, src) = do+  result <- try $ do+    tree <- case parseSR (raFormat args) B.empty True src of+      Left e  -> fail ("parse error: " ++ e)+      Right t -> return $! relabelParams t+    let dist  = raDist args+        nRows = U.length ys+        nModelParams = countParamsUniq tree+        nParams = nModelParams+                + case dist of+                    Gaussian  -> 1+                    ROXY      -> 3+                    _         -> 0++    let et     = compileTree dist xss ys mYerr tree+        theta0 = U.replicate nParams 1.0+        thetaOpt = ctOptimizer et theta0++    when (any isNaN (U.toList thetaOpt)) $+         fail "optimisation returned NaN"++    let mseTree  = buildLoss MSE (fromIntegral nRows) tree+        mseLoss  = compileLoss xss mseTree ys mYerr thetaOpt+        nllLoss  = ctNLL et thetaOpt+        tss      = ctVar et++    let fisherDiag = fisherNLL dist mYerr xss ys tree thetaOpt+        hessCols   = hessianNLL dist mYerr xss ys tree thetaOpt+        hessLists  = map U.toList hessCols+        logP       = logParameters fisherDiag thetaOpt+        logPLatt   = logParametersLatt hessLists fisherDiag thetaOpt+        logF       = logFunctional tree+        logFFreq   = logFunctionalFreq tree+        nF         = fromIntegral nRows+        kF         = fromIntegral nParams+        crits      = map (\c -> (c, evalOne c mseLoss nllLoss tss nF kF logP logPLatt logF logFFreq))+                         (raCriteria args)++    let stats = getStatsFromModel dist mYerr xss ys tree thetaOpt+        laplaceCIs = paramCI (Laplace stats) nRows thetaOpt (raAlpha args)+    let ptype = case raCIType args of+          BatesArg       -> Bates+          ODEArg         -> ODE+          ConstrainedArg -> Constrained+    let kInt = U.length thetaOpt+        nInt = U.length ys+        profT = sqrt $ quantile (fDistribution (fromIntegral kInt) (fromIntegral $ nInt - kInt)) (1 - raAlpha args)+    cis <- case raCI args of+      LaplaceCI -> return laplaceCIs+      ProfileCI -> do+        let profiles = getAllProfiles ptype et thetaOpt (_stdErr stats) laplaceCIs (raAlpha args)+        when (raDbg args) $ forM_ (zip [0..] profiles) $ \(i, ProfileT taus thetas _ tau2theta _) -> do+          putStrLn $ "DEBUG Profile " ++ show i ++ " (opt=" ++ show (thetaOpt U.! i) ++ "):"+          putStrLn $ "  tau range: [" ++ show (if U.null taus then 0 else U.head taus)+                   ++ ", " ++ show (if U.null taus then 0 else U.last taus) ++ "]"+          putStrLn $ "  t=" ++ show profT+          putStrLn $ "  tau2theta(-t)=" ++ show (tau2theta (-profT))+                   ++ "  tau2theta(+t)=" ++ show (tau2theta profT)+          putStrLn $ "  profile points:"+          let tausL = U.toList taus+              thetasL = U.toList (getCol i thetas)+          forM_ (zip tausL thetasL) $ \(tau, th) ->+            putStrLn $ "    tau=" ++ show tau ++ "  theta=" ++ show th+        return $ paramCI (Profile stats profiles) nRows thetaOpt (raAlpha args)++    return $! ReportData+      { rdTree     = tree+      , rdTheta    = thetaOpt+      , rdStdErr   = _stdErr stats+      , rdCriteria = crits+      , rdCIs      = cis+      }++  case result of+    Right rd -> printReport idx src rd+    Left  e  -> printFailure idx src (show (e :: SomeException))++----------------------------------------------------------------------+-- Evaluate a single ModelEval from base quantities+----------------------------------------------------------------------+evalOne :: ModelEval -> Double -> Double -> Double -> Double -> Double+        -> Double -> Double -> Double -> Double -> Double+evalOne RMSE     mse _   _   _ _ _ _ _ _ = sqrt mse+evalOne R2       mse _   tss n _ _ _ _ _ = 1 - n * mse / tss+evalOne AIC      _   nll _   _ k _ _ _ _ = 2*k + 2*nll+evalOne BIC      _   nll _   n k _ _ _ _ = k * log n + 2*nll+evalOne Evidence _   nll _   n k _ _ _ _ = (1 - b) * nll - k/2 * log b+  where b = 1 / sqrt n+evalOne FBF      _   nll _   n k _ _ _ _ = res+  where b = 1 / sqrt n; nup = exp (1 - log 3)+        res = (1 - b) * nll - k/2 * log b + k/2 * log (2*pi*nup)+evalOne MDL      _   nll _   _ _ logP  _ logF _    = nll + logF + logP+evalOne MDLLatt  _   nll _   _ _ _     logPL logF _ = nll + logF + logPL+evalOne MDLFreq  _   nll _   _ _ logP  _ _    logFF = nll + logFF + logP+evalOne (EvalLoss (NLL Gaussian))  _   nll _   _ _ _  _ _    _ = nll+evalOne _        _   _   _   _ _ _     _    _    _  = 0  -- unreachable++----------------------------------------------------------------------+-- Output+----------------------------------------------------------------------+printReport :: Int -> B.ByteString -> ReportData -> IO ()+printReport idx src rd = do+  putStrLn $ "=== Expression " ++ show idx ++ " ==="+  putStrLn $ "Tree: " ++ showExpr (rdTree rd)+  putStrLn "Parameters:"+  let thetaList = U.toList (rdTheta rd)+      ciList    = rdCIs rd+  forM_ (zip3 [0..] thetaList ciList) $ \(i, th, ci) ->+    putStrLn $ "  theta" ++ show i ++ ": " ++ fmt th+            ++ " [" ++ fmt (lower_ ci) ++ ", " ++ fmt (upper_ ci) ++ "]"+  putStrLn "Model Selection:"+  forM_ (rdCriteria rd) $ \(c, v) ->+    putStrLn $ "  " ++ padRight 12 (show c) ++ ": " ++ fmt v+  putStrLn ""+  where+    fmt x | abs x < 1e-10 = "0.0000"+          | abs x >= 1e4  = printf "%.4e" x+          | otherwise     = printf "%.6f" x+    padRight n s = s ++ replicate (max 0 (n - length s)) ' '++printFailure :: Int -> B.ByteString -> String -> IO ()+printFailure idx src msg = do+  putStrLn $ "=== Expression " ++ show idx ++ " ==="+  putStrLn $ "Tree: " ++ B.unpack src+  putStrLn $ "Error: " ++ msg+  putStrLn ""

@@ -1,103 +0,0 @@-module Main (main) where--import Options.Applicative-import Text.ParseSR.IO ( withInput, withOutput )-import Text.ParseSR ( SRAlgs (..), Output (..) )-import System.Random ( getStdGen, mkStdGen )-import Text.Read ( readMaybe )-import Data.Char ( toLower, toUpper )-import Data.List ( intercalate )---- Data type to store command line arguments-data Args = Args-    {   from        :: SRAlgs-      , to          :: Output-      , infile      :: String-      , outfile     :: String-      , varnames    :: String-    } deriving Show---- parser of command line arguments-opt :: Parser Args-opt = Args-   <$> option sralgsReader-       ( long "from"-       <> short 'f'-       <> metavar ("[" <> intercalate "|" sralgsHelp <> "]")-       <> help "Input expression format" )-   <*> option srtoReader -- TODO-       ( long "to"-       <> short 't'-       <> metavar ("[" <> intercalate "|" srtoHelp <> "]")-       <> help "Output expression format" )-   <*> strOption-       ( long "input"-       <> short 'i'-       <> metavar "INPUT-FILE"-       <> showDefault-       <> value ""-       <> help "Input file containing expressions. \-               \ Empty string gets expression from stdin." )-   <*> strOption-       ( long "output"-       <> short 'o'-       <> metavar "OUTPUT-FILE"-       <> showDefault-       <> value ""-       <> help "Output file to store the stats in CSV format. \-                \ Empty string prints expressions to stdout." )-   <*> strOption-      ( long "varnames"-      <> short 'v'-      <> metavar "VARNAMES"-      <> showDefault-      <> value ""-      <> help "Comma separated string of variable names. \-               \ Empty string defaults to the algorithm default (x0, x1,..)." )---- helper functions to show the possible options-mkDescription :: Show a => [a] -> [String]-mkDescription = map (envelope '\'' . map toLower . show)-  where-    envelope :: a -> [a] -> [a]-    envelope c xs = c : xs <> [c]-{-# INLINE mkDescription #-}--sralgsHelp :: [String]-sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..]-{-# INLINE sralgsHelp #-}--srtoHelp :: [String]-srtoHelp = mkDescription [toEnum 0 :: Output ..]-{-# INLINE srtoHelp #-}---- helper functions to parse the options-mkReader :: Read a => String -> (a -> b) -> String -> ReadM b-mkReader err val sr = eitherReader-                    $ case readMaybe sr of-                        Nothing -> pure (Left err)-                        Just x  -> pure (Right (val x))--sralgsReader :: ReadM SRAlgs-sralgsReader =-  str >>= (mkReader errMsg id . map toUpper)-  where-    errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp--srtoReader :: ReadM Output-srtoReader =-  str >>= (mkReader errMsg id . map toUpper)-  where-    errMsg = "unknown algorithm. Available options are " <> intercalate "," srtoHelp--main :: IO ()-main = do-  args <- execParser opts-  withInput (infile args) (from args) (varnames args) False True-    >>= withOutput (outfile args) (to args)-  where-    opts = info (opt <**> helper)-            ( fullDesc <> progDesc "Simplify an expression\-                                   \ using equality saturation."-           <> header "srsimplify - a CLI tool to simplify\-                     \ symbolic regression expressions with equality saturation." )

@@ -1,184 +0,0 @@-module Args where--import Data.Char ( toLower, toUpper )-import Data.List ( intercalate )-import Algorithm.SRTree.Likelihoods ( Distribution (..) )-import Algorithm.SRTree.ConfidenceIntervals ( PType (..) )-import Options.Applicative-import Text.ParseSR ( SRAlgs (..) )-import Text.Read ( readMaybe )---- Data type to store command line arguments-data Args = Args-    {   from        :: SRAlgs-      , infile      :: String-      , outfile     :: String-      , dataset     :: String-      , test        :: String-      , niter       :: Int-      , hasHeader   :: Bool-      , simpl       :: Bool-      , dist        :: Distribution-      , restart     :: Bool-      , rseed       :: Int-      , toScreen    :: Bool-      , useProfile  :: Bool-      , simple      :: Bool-      , sigma       :: Double-      , alpha       :: Double-      , ptype       :: PType-    } deriving Show---- parser of command line arguments-opt :: Parser Args-opt = Args-   <$> option sralgsReader-       ( long "from"-       <> short 'f'-       <> metavar ("[" <> intercalate "|" sralgsHelp <> "]")-       <> help "Input expression format" )-   <*> strOption-       ( long "input"-       <> short 'i'-       <> metavar "INPUT-FILE"-       <> showDefault-       <> value ""-       <> help "Input file containing expressions. \-               \ Empty string gets expression from stdin." )-   <*> strOption-       ( long "output"-       <> short 'o'-       <> metavar "OUTPUT-FILE"-       <> showDefault-       <> value ""-       <> help "Output file to store the stats in CSV format. \-                \ Empty string prints expressions to stdout." )-   <*> strOption-       ( long "dataset"-       <> short 'd'-       <> metavar "DATASET-FILENAME"-       <> help "Filename of the dataset used for optimizing the parameters. \-               \ Empty string omits stats that make use of the training data. \-               \ It will auto-detect and handle gzipped file based on gz extension. \-               \ It will also auto-detect the delimiter.\n\-               \ The filename can include extra information: \-               \ filename.csv:start:end:target:cols:yerr:xerr where start and end \-               \ corresponds to the range of rows that should be used for fitting,\-               \ target is the column index (or name) of the target variable and cols\-               \ is a comma separated list of column indices or names of the variables\-               \ in the same order as used by the symbolic model.\-               \ The yerr field corresponds to the column with the error of the target,\-               \ while xerr a comma separated indices of the columns with the error of the\-               \ variables. If nothing passed, it will ignore measurement errors." )-   <*> strOption-       ( long "test"-       <> metavar "TEST"-       <> showDefault-       <> value ""-       <> help "Filename of the test dataset.\-               \ Empty string omits stats that make use of the training data.\-               \ It can have additional information as in the training set,\-               \ but the validation range will be discarded." )-   <*> option auto-       ( long "niter"-       <> metavar "NITER"-       <> showDefault-       <> value 10-       <> help "Number of iterations for the optimization algorithm.")-   <*> switch-       ( long "hasheader"-       <> help "Uses the first row of the csv file as header.")-   <*> switch-        ( long "simplify"-        <> help "Apply basic simplification." )-   <*> option distRead-        ( long "distribution"-        <> metavar ("[" <> intercalate "|" distHelp <> "]")-        <> showDefault-        <> value Gaussian-        <> help "Minimize negative log-likelihood following one of\-                \ the avaliable distributions.\-                \ The default is Gaussian."-        )-   <*> switch-        ( long "restart"-        <> help "If set, it samples the initial values of\-                 \ the parameters using a Gaussian distribution N(0, 1),\-                 \ otherwise it uses the original values of the expression." )-   <*> option auto-       ( long "seed"-       <> metavar "SEED"-       <> showDefault-       <> value (-1)-       <> help "Random seed to initialize the parameters values.\-                \ Used only if restart is enabled.")-   <*> switch-        ( long "report"-        <> help "If set, reports the analysis in a user-friendly\-                \ format instead of csv. It will also include\-                \ confidence interval for the parameters and predictions" )-   <*> switch-        ( long "profile"-        <> help "If set, it will use profile likelihood to calculate the CIs." )-   <*> switch-       ( long "simple"-       <> help "If set, calculates only SSE.")-   <*> option auto-       ( long "sigma"-       <> metavar "SIGMA"-       <> showDefault-       <> value 0.001-       <> help "Estimation of error for Guassian distribution.")-   <*> option auto-       ( long "alpha"-       <> metavar "ALPHA"-       <> showDefault-       <> value 0.05-       <> help "Significance level for confidence intervals.")-    <*> option auto-        ( long "ptype"-        <> metavar "[Bates | ODE | Constrained]"-        <> showDefault-        <> value Constrained-        <> help "Profile Likelihood method. Default: Constrained. NOTE: Constrained method only calculates the endpoint."-        )---- helper functions to show the possible options-mkDescription :: Show a => [a] -> [String]-mkDescription = map (envelope '\'' . map toLower . show) -  where-    envelope :: a -> [a] -> [a]-    envelope c xs = c : xs <> [c]-{-# INLINE mkDescription #-}--sralgsHelp :: [String]-sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..]-{-# INLINE sralgsHelp #-}--distHelp :: [String]-distHelp = mkDescription [toEnum 0 :: Distribution ..]-{-# INLINE distHelp #-}---- helper functions to parse the options-mkReader :: Read a => String -> (a -> b) -> String -> ReadM b-mkReader err val sr = eitherReader -                    $ case readMaybe sr of-                        Nothing -> pure (Left err)-                        Just x  -> pure (Right (val x))--sralgsReader :: ReadM SRAlgs-sralgsReader =-  str >>= (mkReader errMsg id . map toUpper)-  where-    errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp----s2Reader :: ReadM (Maybe Double)---s2Reader =---  str >>= \s -> mkReader ("wrong format " <> s) Just s--distRead :: ReadM Distribution-distRead =-  str >>= \s -> mkReader ("unsupported distribution " <> s) id (capitalize s)-  where-    capitalize ""     = ""-    capitalize (c:cs) = toUpper c : if length cs == 2 then map toUpper cs else map toLower cs

@@ -1,230 +0,0 @@-{-# language BlockArguments #-}-{-# language LambdaCase #-}-module IO where--import System.IO ( hClose, hPutStrLn, openFile, stderr, stdout, IOMode(WriteMode), Handle )-import qualified Data.Massiv.Array as A-import Data.List ( intercalate )-import Control.Monad ( unless, forM_ )-import System.Random ( StdGen )--import Data.SRTree ( SRTree (..), Fix (..), var, floatConstsToParam, relabelVars )-import Algorithm.SRTree.Likelihoods ( Distribution (..) )-import Algorithm.SRTree.ConfidenceIntervals ( printCI, BasicStats(_stdErr, _corr), CI )-import qualified Data.SRTree.Print as P-import Data.SRTree.Eval ( compMode )--import Args ( Args(outfile, alpha,dist,niter,sigma) )-import Report-import Data.SRTree.Recursion ( cata )--import Debug.Trace ( trace, traceShow )---- Header of CSV file-csvHeader :: String-csvHeader = intercalate "," (basicFields <> optFields <> modelFields)-{-# inline csvHeader #-}--csvHeaderSimple :: String-csvHeaderSimple = intercalate "," (basicFields <> optFields)-{-# inline csvHeaderSimple #-}---- Open file if filename is not empty-openWriteWithDefault :: Handle -> String -> IO Handle-openWriteWithDefault dflt ""    = pure dflt-openWriteWithDefault _    fname = openFile fname WriteMode-{-# INLINE openWriteWithDefault #-}---- procecss a single tree and return all the available stats-processTree :: Args        -- command line arguments-            -> StdGen      -- random number generator-            -> Datasets    -- datasets-            -> Fix SRTree  -- expression in tree format-            -> Int         -- index of the parsed expression -            -> (BasicInfo, SSE, SSE, Info, (BasicStats, [CI], [CI], [CI], [CI]))-processTree args seed dset t ix = (basic, sseOrig, sseOpt, info, cis)-  where-    (tree, theta0')  = floatConstsToParam t-    theta0           = if dist args == Gaussian-                          then theta0' <> [sigma args]-                          else theta0'--    basic   = getBasicStats args seed dset tree theta0 ix-    treeVal = case (_xVal dset, _yVal dset) of-                (Nothing, _) -> _expr basic-                (_, Nothing) -> _expr basic-                (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix-    sseOrig = getSSE dset t-    sseOpt  = getSSE dset (_expr basic)-    info    = getInfo args dset (_expr basic) treeVal-    cis     = getCI args dset basic (alpha args)--processTreeSimple :: Args        -- command line arguments-            -> StdGen      -- random number generator-            -> Datasets    -- datasets-            -> Fix SRTree  -- expression in tree format-            -> Int         -- index of the parsed expression-            -> (BasicInfo, SSE, SSE)-processTreeSimple args seed dset t ix = (basic, sseOrig, sseOpt)-  where-    (tree, theta0')  = floatConstsToParam t-    theta0           = if dist args == Gaussian-                          then theta0' <> [sigma args]-                          else theta0'--    basic   = getBasicStats args seed dset tree theta0 ix-    treeVal = case (_xVal dset, _yVal dset) of-                (Nothing, _) -> _expr basic-                (_, Nothing) -> _expr basic-                (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix-    sseOrig = getSSE dset t-    sseOpt  = getSSE dset (_expr basic)---- print the results to a csv format (except CI)-printResults :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO ()-printResults args seed dset varnames exprs  = do-  hStat <- openWriteWithDefault stdout (outfile args)-  hPutStrLn hStat csvHeader -  forM_ (zip [0..] exprs) -     \(ix, tree) -> -         case tree of-           Left  err -> hPutStrLn stderr ("invalid expression: " <> err)-           Right t   -> let treeData = processTree args seed dset t ix-                        in hPutStrLn hStat (toCsv treeData varnames)-  unless (null (outfile args)) (hClose hStat)--printResultsSimple :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO ()-printResultsSimple args seed dset varnames exprs  = do-  hStat <- openWriteWithDefault stdout (outfile args)-  hPutStrLn hStat csvHeaderSimple-  forM_ (zip [0..] exprs)-     \(ix, tree) ->-         case tree of-           Left  err -> hPutStrLn stderr ("invalid expression: " <> err)-           Right t   -> let treeData = processTreeSimple args seed dset t ix-                        in hPutStrLn hStat (toCsvSimple treeData varnames)-  unless (null (outfile args)) (hClose hStat)---- change the stats into a string-toCsv :: (BasicInfo, SSE, SSE, Info, e) -> [String] -> String-toCsv (basic, sseOrig, sseOpt, info, _) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt <> sInfo)-  where-    sBasic    = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic)-                , show (_nNodes basic), show (_nParams basic)-                , intercalate ";" (map show (_params basic))-                , show (_nEvals basic)-                ]-    sSSEOrig  = map (showF sseOrig) [_sseTr, _sseVal, _sseTe]-    sSSEOpt   = map (showF sseOpt)  [_sseTr, _sseVal, _sseTe]-    sInfo     = map (showF info) [_bic, _bicVal, _aic, _aicVal, _evidence, _evidenceVal, _mdl, _mdlFreq, _mdlLatt, _mdlVal, _mdlFreqVal, _mdlLattVal, _nllTr, _nllVal, _nllTe, _cc, _cp]-              <> [intercalate ";" (map show (_fisher info))]-    showF p f = show (f p)--toCsvSimple :: (BasicInfo, SSE, SSE) -> [String] -> String-toCsvSimple (basic, sseOrig, sseOpt) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt)-  where-    sBasic    = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic)-                , show (_nNodes basic), show (_nParams basic)-                , intercalate ";" (map show (_params basic))-                , show (_nEvals basic)-                ]-    sSSEOrig  = map (showF sseOrig) [_sseTr, _sseVal, _sseTe]-    sSSEOpt   = map (showF sseOpt)  [_sseTr, _sseVal, _sseTe]-    showF p f = show (f p)---- get trees of transformed features-getTransformedFeatures :: Fix SRTree -> (Fix SRTree, [Fix SRTree])-getTransformedFeatures = cata $-  \case-     Var   ix                   -> (Fix $ Var ix, [])-     Param ix                   -> (Fix $ Param ix, [])-     Const  x                   -> (Fix $ Const x, [])-     Uni    f (t, vars)         -> (Fix $ Uni f t, vars)-     Bin   op (l, vs1) (r, vs2) -> case (hasNoParam l, hasNoParam r) of-                                     (False, True)  -> let vs = vs1 <> vs2-                                                       in (Fix $ Bin op l (var $ length vs), vs <> [r])-                                     (True, False)  -> let vs = vs1 <> vs2-                                                       in (Fix $ Bin op (var $ length vs) r, vs <> [l])-                                     (    _,    _)   -> (Fix $ Bin op l r, vs1 <> vs2) -- vs1 == vs2 == []-- where-   hasNoParam = cata $-     \case-        Var ix     -> True-        Param ix   -> False-        Const x    -> if floor x == ceiling x then True else False-        Uni f t    -> t-        Bin op l r -> l && r--allAreVars :: [Fix SRTree] -> Bool-allAreVars = all isOnlyVar-  where-    isOnlyVar (Fix (Var _)) = True-    isOnlyVar _             = False---- print the information on screen (including CIs)-printResultsScreen :: Args -> StdGen -> Datasets -> [String] -> String -> [Either String (Fix SRTree)] -> IO ()-printResultsScreen args seed dset varnames targt exprs = do-  forM_ (zip [0..] exprs) -    \(ix, tree) -> -        case tree of-          Left  err -> do putStrLn ("invalid expression: " <> err)-          Right t   -> let treeData = processTree args seed dset t ix-                        in printToScreen ix treeData-  where-    decim :: Int -> Double -> Double-    decim n x = (fromIntegral . (round :: Double -> Integer)) (x * 10^n) / 10^n-    sdecim n  = show . decim n-    nplaces   = 4---    printToScreen ix (basic, _, sseOpt, info, (sts, cis, pis_tr, pis_val, pis_te)) =-      do let (transformedT, newvars) = getTransformedFeatures (_expr basic)-             varnames' = ['z': show ix | ix <- [0 .. length newvars - 1]]-         putStrLn $ "=================== EXPR " <> show ix <> " =================="-         putStr   $ targt <> " ~ f(" <> intercalate ", " varnames <> ") = "-         putStrLn $ P.showExprWithVars varnames (_expr basic)--         unless (allAreVars newvars) do-          putStrLn "\nExpression and transformed features: "-          putStr   $ targt <> " ~ f(" <> intercalate ", " varnames' <> ") = "-          putStrLn $ P.showExprWithVars varnames' (relabelVars transformedT)-          forM_ (zip varnames' newvars) \(vn, tv) -> do-            putStrLn $ vn <> " = " <> P.showExprWithVars varnames tv--         putStrLn "\n---------General stats:---------\n"-         putStrLn $ "Number of nodes: " <> show (_nNodes basic)-         putStrLn $ "Number of params: " <> show (_nParams basic)-         putStrLn $ "theta = " <> show (_params basic)--         putStrLn "\n----------Performance:--------\n"-         putStrLn $ "SSE (train.): " <> sdecim nplaces (_sseTr sseOpt)-         putStrLn $ "SSE (val.): " <> sdecim nplaces (_sseVal sseOpt)-         putStrLn $ "SSE (test): " <> sdecim nplaces (_sseTe sseOpt)-         putStrLn $ "NegLogLiklihood (train.): " <> sdecim nplaces (_nllTr info)-         putStrLn $ "NegLogLiklihood (val.): " <> sdecim nplaces (_nllVal info)-         putStrLn $ "NegLogLiklihood (test): " <> sdecim nplaces (_nllTe info)--         putStrLn "\n------Selection criteria:-----\n"-         putStrLn $ "BIC: " <> sdecim nplaces (_bic info)-         putStrLn $ "AIC: " <> sdecim nplaces (_aic info)-         putStrLn $ "MDL: " <> sdecim nplaces (_mdl info)-         putStrLn $ "MDL (freq.): " <> sdecim nplaces (_mdlFreq info)-         putStrLn $ "Functional complexity: " <> sdecim nplaces (_cc info)-         putStrLn $ "Parameter complexity: " <> sdecim nplaces (_cp info)--         putStrLn "\n---------Uncertainties:----------\n"-         putStrLn "Correlation of parameters: " -         putStrLn $ show $ A.map (decim 2) (_corr sts)-         putStrLn $ "Std. Err.: " <> show (A.map (decim nplaces) (_stdErr sts))-         putStrLn "\nConfidence intervals:\n\nlower <= val <= upper"-         mapM_ (printCI nplaces) cis-         putStrLn "\nConfidence intervals (predictions training):\n\nlower <= val <= upper"-         mapM_ (printCI nplaces) pis_tr-         unless (null pis_val) do-           putStrLn "\nConfidence intervals (predictions validation):\n\nlower <= val <= upper"-           mapM_ (printCI nplaces) pis_val-         unless (null pis_te) do-           putStrLn "\nConfidence intervals (predictions test):\n\nlower <= val <= upper"-           mapM_ (printCI nplaces) pis_te-         putStrLn "============================================================="

@@ -1,34 +0,0 @@-module Main (main) where--import Data.ByteString.Char8 ( pack, unpack, split )-import Options.Applicative-import System.Random ( getStdGen, mkStdGen )-import Text.ParseSR.IO ( withInput )--import Args-import IO-import Report--main :: IO ()-main = do-  args             <- execParser opts-  g                <- getStdGen-  (dset, varnames, tgname) <- getDataset args--  let seed = if rseed args < 0 -               then g -               else mkStdGen (rseed args)-      varnames' = map unpack $ split ',' $ pack varnames-  withInput (infile args) (from args) varnames False (simpl args)-    >>= if toScreen args-          then printResultsScreen args seed dset varnames' tgname  -- full report on screen-          else if simple args-                 then printResultsSimple args seed dset varnames' -- csv file-                 else printResults args seed dset varnames' -- csv file-  where    -    opts = info (opt <**> helper)-            ( fullDesc <> progDesc "Optimize the parameters of\-                                   \ Symbolic Regression expressions."-           <> header "srtools - a CLI tool to (re)optimize the numeric\-                     \ parameters of symbolic regression expressions"-            )

@@ -1,280 +0,0 @@-module Report where--import qualified Data.Vector.Storable as VS-import qualified Data.Massiv.Array as A-import Data.Massiv.Array ( Sz(..) )-import Data.Maybe ( fromMaybe )-import Statistics.Distribution.FDistribution ( fDistribution )-import Statistics.Distribution.ChiSquared ( chiSquared )-import Statistics.Distribution ( quantile )-import System.Random ( StdGen, split, randomRs )--import Data.SRTree ( SRTree, Fix (..), floatConstsToParam, paramsToConst, countNodes )-import Data.SRTree.Eval-import Algorithm.SRTree.AD ( forwardModeUniqueJac )-import Algorithm.SRTree.Likelihoods-import Algorithm.SRTree.ModelSelection ( aic, bic, evidence, logFunctional, logParameters, mdl, mdlFreq, mdlLatt )-import Algorithm.SRTree.ConfidenceIntervals-import Algorithm.SRTree.Opt (minimizeNLLWithFixedParam, minimizeNLL)-import Data.SRTree.Datasets ( loadDataset )-import Data.SRTree.Print ( showExpr )-import Debug.Trace ( trace, traceShow )--import Args---- store the datasets split into training, validation and test-data Datasets = DS { _xTr     :: SRMatrix-                   , _yTr     :: PVector-                   , _xVal    :: Maybe SRMatrix-                   , _yVal    :: Maybe PVector-                   , _xTe     :: Maybe SRMatrix-                   , _yTe     :: Maybe PVector-                   , _yErrTr  :: Maybe PVector-                   , _yErrVal :: Maybe PVector-                   , _yErrTe  :: Maybe PVector-                   }---- basic fields name-basicFields :: [String]-basicFields = [ "Index"-              , "Filename"-              , "Expression"-              , "Number_of_nodes"-              , "Number_of_parameters"-              , "Parameters"-              , "Number_of_evaluations"-              ]---- basic information about the tree-data BasicInfo = Basic { _index   :: Int-                       , _fname   :: String-                       , _expr    :: Fix SRTree-                       , _nNodes  :: Int-                       , _nParams :: Int-                       , _params  :: [Double]-                       , _nEvals  :: Int-                       }---- optimization fields-optFields :: [String]-optFields = [ "SSE_train_orig"-            , "SSE_val_orig"-            , "SSE_test_orig"-            , "SSE_train_opt"-            , "SSE_val_opt"-            , "SSE_test_opt"-            ]---- optimization information-data SSE = SSE { _sseTr  :: Double-               , _sseVal :: Double-               , _sseTe  :: Double-               }---- model selection fields-modelFields :: [String]-modelFields = [ "BIC"-              , "BIC_val"-              , "AIC"-              , "AIC_val"-              , "Evidence"-              , "EvidenceVal"-              , "MDL"-              , "MDL_Freq"-              , "MDL_Lattice"-              , "MDL_val"-              , "MDL_Freq_val"-              , "MDL_Lattice_val"-              , "NegLogLikelihood_train"-              , "NegLogLikelihood_val"-              , "NegLogLikelihood_test"-              , "LogFunctional"-              , "LogParameters"-              , "Fisher"-              ]---- model selection information-data Info = Info { _bic     :: Double-                 , _bicVal  :: Double-                 , _aic     :: Double-                 , _aicVal  :: Double-                 , _evidence :: Double-                 , _evidenceVal :: Double-                 , _mdl     :: Double-                 , _mdlFreq :: Double-                 , _mdlLatt :: Double-                 , _mdlVal  :: Double-                 , _mdlFreqVal :: Double-                 , _mdlLattVal :: Double-                 , _nllTr   :: Double-                 , _nllVal  :: Double-                 , _nllTe   :: Double-                 , _cc      :: Double-                 , _cp      :: Double-                 , _fisher  :: [Double]-                 }---- load the datasets-getDataset :: Args -> IO (Datasets, String, String)-getDataset args = do-  ((xTr, yTr, xVal, yVal), (yErrTr, yErrVal), varnames, tgname) <- loadDataset (dataset args) (hasHeader args)-  let (A.Sz m) = A.size yVal-  let (mXVal, mYVal) = if m == 0-                         then (Nothing, Nothing)-                         else (Just xVal, Just yVal)-  (mXTe, mYTe, mYErrTe) <- if null (test args)-                             then pure (Nothing, Nothing, Nothing)-                             else do ((xTe, yTe, _, _), (yErrTe, _), _, _) <- loadDataset (test args) (hasHeader args)-                                     pure (Just xTe, Just yTe, yErrTe)-  pure (DS xTr yTr mXVal mYVal mXTe mYTe yErrTr yErrVal mYErrTe, varnames, tgname)--getBasicStats :: Args -> StdGen -> Datasets -> Fix SRTree -> [Double] -> Int -> BasicInfo-getBasicStats args seed dset tree theta0 ix-  | anyNaN    = getBasicStats args (snd $ split seed) dset tree theta0 ix-  | otherwise = Basic ix (infile args) tOpt nNodes nParams params nEvs-  where-    -- (tree', theta0) = floatConstsToParam tree-    thetas          = if restart args-                        then A.fromList compMode $ take nParams (randomRs (-1.0, 1.0) seed)-                        else A.fromList compMode theta0-    (t,_,nEvs)      = minimizeNLL (dist args) (_yErrTr dset) (niter args) (_xTr dset) (_yTr dset) tree thetas-    tOpt            = paramsToConst (A.toList t) tree-    nNodes          = countNodes tOpt :: Int-    nParams         = length theta0-    params          = A.toList t-    anyNaN          = A.any isNaN t--getSSE :: Datasets -> Fix SRTree -> SSE-getSSE dset tree = SSE tr val te-  where-    (t, th) = floatConstsToParam tree-    tr  = sse (_xTr dset) (_yTr dset) t (A.fromList compMode th)-    val = case (_xVal dset, _yVal dset) of-            (Nothing, _)           -> 0.0-            (_, Nothing)           -> 0.0-            (Just xVal, Just yVal) -> sse xVal yVal t (A.fromList compMode th)-    te  = case (_xTe dset, _yTe dset) of-            (Nothing, _)           -> 0.0-            (_, Nothing)           -> 0.0-            (Just xTe, Just yTe)   -> sse xTe yTe t (A.fromList compMode th)--getInfo :: Args -> Datasets -> Fix SRTree -> Fix SRTree -> Info-getInfo args dset tree treeVal =-  Info { _bic     = bic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _bicVal  = bicVal-       , _aic     = aic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _aicVal  = aicVal-       , _evidence = evidence dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _evidenceVal = evidenceVal-       , _mdl     = mdl dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _mdlFreq = mdlFreq dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _mdlLatt = mdlLatt dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _mdlVal  = mdlVal-       , _mdlFreqVal = mdlFreqVal-       , _mdlLattVal = mdlLattVal-       , _nllTr   = nllTr-       , _nllVal  = nllVal-       , _nllTe   = nllTe-       , _cc      = logFunctional tOpt-       , _cp      = logParameters dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt-       , _fisher  = A.toList $ fisherNLL dist' (_yErrTr dset) xTr yTr tOpt thetaOpt'-       }-  where-    (xTr, yTr)       = (_xTr dset, _yTr dset)-    (xVal, yVal)     = case (_xVal dset, _yVal dset) of-                         (Nothing, _)     -> (xTr, yTr)-                         (_, Nothing)     -> (xTr, yTr)-                         (Just a, Just b) -> (a, b)-    (tOpt, thetaOpt_nosig) = floatConstsToParam tree-    thetaOpt         = if dist args == Gaussian-                          then thetaOpt_nosig <> [sigma args]-                          else thetaOpt_nosig-    thetaOpt'        = A.fromList compMode thetaOpt--    (tOptVal, thetaOptVal_nosig) = floatConstsToParam treeVal-    thetaOptVal  = if dist args == Gaussian-                      then thetaOptVal_nosig <> [sigma args]-                      else thetaOptVal_nosig-    thetaOptVal'           = A.fromList compMode thetaOptVal--    dist'            = dist args--    nllTr            = nll dist' (_yErrTr dset) (_xTr dset) (_yTr dset) tOpt (A.fromList compMode thetaOpt)-    bicVal           = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> bic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    aicVal           = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> aic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    evidenceVal      = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> evidence dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    nllVal           = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> nll dist' (_yErrVal dset) xVal yVal tOptVal (A.fromList compMode thetaOptVal)-    mdlVal           = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> mdl dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    mdlFreqVal       = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> mdlFreq dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    mdlLattVal       = case (_xVal dset, _yVal dset) of-                         (Nothing, _) -> 0.0-                         (_, Nothing) -> 0.0-                         _            -> mdlLatt dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal-    nllTe            = case (_xTe dset, _yTe dset) of-                         (Nothing, _)           -> 0.0-                         (_, Nothing)           -> 0.0-                         (Just xTe, Just yTe) -> nll dist' (_yErrTe dset) xTe yTe tOpt (A.fromList compMode thetaOpt)--getCI :: Args -> Datasets -> BasicInfo -> Double -> (BasicStats, [CI], [CI], [CI], [CI])-getCI args dset basic alpha' = (stats', cis, pis_tr, pis_val, pis_te)-  where-    (Sz n)     = A.size yTr-    (tree, _)  = floatConstsToParam (_expr basic)-    theta      = _params basic-    tau_max    = (quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - 0.01))-    tau_max'   = sqrt $ quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - alpha')-    (xTr, yTr) = (_xTr dset, _yTr dset)-    dist'      = dist args-    stats'     = getStatsFromModel dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta)-    profiles   = getAllProfiles (ptype args) dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta) (_stdErr stats') estCIs alpha'-    method     = if useProfile args-                   then Profile stats' profiles-                   else Laplace stats'-    predFun   = A.computeAs A.S . predict dist' tree (A.fromList compMode theta)--    prof estPi th t =-                let (thOpt, _, _) = minimizeNLL dist' (_yErrTr dset) 100 xTr yTr t th-                    ssr        = sse xTr yTr t thOpt-                    est        = sqrt $ ssr / fromIntegral (n - _nParams basic)-                    stdErr     = _stdErr stats' A.! 0-                    fun        = case ptype args of-                                   Bates       -> getProfile      dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max 0-                                   ODE         -> getProfileODE   dist' (_yErrTr dset) xTr yTr t thOpt stdErr estPi tau_max 0-                                   Constrained -> getProfileCnstr dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max' 0-                in case fun of-                      Left th' -> trace "found better optima" $ prof estPi th' t-                      Right p  -> (_tau2theta p, _opt p)-    jac xss   = forwardModeUniqueJac xss (A.fromList compMode theta) tree -- FIX--    estCIs    = paramCI (Laplace stats') n (A.fromList compMode theta) 0.001-    cis       = paramCI method n (A.fromList compMode theta) alpha'--    estPIS_tr  = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []-    estPIS_val = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []-    estPIS_te  = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []--    pis_tr    = predictionCI method dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' estPIS_tr-    pis_val   = case (_xVal dset, _yVal dset) of-                  (Nothing, _)   -> []-                  (Just xVal, _) -> predictionCI method dist' predFun jac prof xVal tree (A.fromList compMode theta) alpha' estPIS_val-    pis_te    = case (_xTe dset, _yTe dset) of-                  (Nothing, _)  -> []-                  (Just xTe, _) -> predictionCI method dist' predFun jac prof xTe tree (A.fromList compMode theta) alpha' estPIS_te

@@ -1,252 +0,0 @@-{-# LANGUAGE ImportQualifiedPost #-}-{-# LANGUAGE TupleSections #-}-{-# LANGUAGE BangPatterns #-}-module GP where--import Data.SRTree-import Algorithm.SRTree.Opt-import Algorithm.SRTree.Likelihoods-import Data.SRTree.Print-import Data.SRTree.Eval-import Data.SRTree.Recursion ( cata )-import System.Random-import Control.Monad.State.Strict-import Control.Monad-import Data.Vector qualified as V-import Control.Monad (when)-import Data.Massiv.Array qualified as M-import Debug.Trace ( traceShow, trace )-import Util-import Data.List ( intercalate, maximumBy )-import qualified Data.Vector.Mutable as MV--data Method = Grow | Full | BTC-type Rng a = StateT StdGen IO a--type GenUni = Fix SRTree -> Fix SRTree -type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree-type FitFun = Individual -> Rng Individual--data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: [PVector] }--instance Show Individual where -    show (Individual t f p) = showExpr t <> "," <> show f <> "," <> show p --toss :: Rng Bool-toss = state random-{-# INLINE toss #-}--randomRange :: (Ord val, Random val) => (val, val) -> Rng val-randomRange rng = state (randomR rng)-{-# INLINE randomRange #-}--randomFrom :: [a] -> Rng a-randomFrom funs = do n <- randomRange (0, length funs - 1)-                     pure $ funs !! n-{-# INLINE randomFrom #-}--randomFromV :: V.Vector a -> Rng a-randomFromV funs = do n <- randomRange (0, length funs - 1)-                      pure $ funs V.! n-{-# INLINE randomFromV #-}--countNodes' :: Fix SRTree -> Int-countNodes' = cata alg -  where -    alg (Var _)     = 1-    alg (Param _)   = 1-    alg (Const _)   = 0-    alg (Bin _ l r) = 1 + l + r-    alg (Uni Abs t) = t-    alg (Uni _ t)   = 1 + t-{-# INLINE countNodes' #-}---randomTree :: HyperParams -> Bool -> Rng (Fix SRTree)-randomTree hp grow -  | depth <= 1 || size <= 2 = randomFrom term -  | (min_depth >= 0 || (depth > 2 && not grow)) && size > 2 = genNonTerm -  | otherwise = genTermOrNon-  where -    min_depth = _minDepth hp-    depth     = _maxDepth hp-    size      = _maxSize hp-    term      = _term hp-    nonterm   = _nonterm hp--    genNonTerm =-       do et <- randomFrom nonterm-          case et of -            Left uniT -> uniT <$> randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow-            Right binT -> do l <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow-                             r <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1 - countNodes' l} grow-                             pure (binT l r)-    genTermOrNon = do r <- toss-                      if r-                        then randomFrom term -                        else genNonTerm        --data HyperParams = -    HP { _minDepth  :: Int -       , _maxDepth  :: Int-       , _maxSize   :: Int -       , _popSize   :: Int-       , _tournSize :: Int-       , _pc        :: Double -       , _pm        :: Double -       , _term      :: [Fix SRTree]-       , _nonterm   :: [Either GenUni GenBin] -       }--tournament :: HyperParams -> V.Vector Individual -> Rng Individual-tournament hp pop = do-  selection <- replicateM (_tournSize hp) (randomFromV $ V.filter (not.isNaN._fit) pop)-  let maxFitness = maximum (fmap _fit selection)-      champions = V.filter ((== maxFitness) . _fit) pop-  if null selection-     then randomFromV pop-     else randomFromV champions--randomIndividual :: HyperParams -> FitFun -> Bool -> Rng Individual-randomIndividual hyperparams fitFun grow = do-    t <- randomTree hyperparams grow -    let p = countParams t-    --theta' <- replicateM p (randomRange (-1,1))-    fitFun $ Individual t 0.0 [] -- (M.fromList compMode theta' :: PVector)-    --pure ind-    --if isInfinite (_fit ind)-    --   then randomIndividual hyperparams fitFun grow -    --   else pure ind--initialPop :: HyperParams -> FitFun -> Rng (V.Vector Individual)-initialPop hyperparams fitFun = do -   let depths = [3 .. _maxDepth hyperparams]-   pop <- forM depths $ \md -> -           do let m = _popSize hyperparams `div` (_maxDepth hyperparams - 3 + 1)-                  g = V.fromList . take m $ cycle [True, False]-              mapM (randomIndividual hyperparams{ _maxDepth = md} fitFun) g-   pure (V.concat pop)--fitnessMV :: Distribution -> [(SRMatrix, PVector, Maybe PVector)] -> Individual -> Rng Individual-fitnessMV dist datas ind = do-  fs <- forM datas (fitness dist ind)-  let fitOpt = minimum $ map fst fs-  pure ind{_fit = fitOpt, _params = map snd fs}--fitness :: Distribution -> Individual ->  (SRMatrix, PVector, Maybe PVector) -> Rng (Double, PVector)-fitness dist ind (x, y, e) = do-    let tree = relabelParams $ _tree ind-        p    = countParams tree-    theta1' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))-    theta2' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))-    let (theta1, f1, _) = minimizeNLL dist e 50 x y tree theta1'-        (theta2, f2, _) = minimizeNLL dist e 50 x y tree theta2'-        fit1 = if isNaN f1 then (-1.0/0.0) else negate f1-        fit2 = if isNaN f1 then (-1.0/0.0) else negate f2-        thetaOpt = if fit1 > fit2 then theta1 else theta2-        fitOpt   = max fit1 fit2-    pure (fitOpt, thetaOpt)---mutate :: HyperParams -> Individual -> Rng (Maybe Individual)-mutate hp ind = do-  let sz = countNodes' (_tree ind)-  p <- state $ randomR (0, sz-1)-  b <- state random-  t <- go p (_maxSize hp) (_tree ind)-  --(t, b) <- go sz (_pm hp) (_tree ind)-  if b <= _pm hp && countNodes t <= _maxSize hp-     then pure . Just $ Individual t 0.0 []-     else pure Nothing-      where-        go 0 msz t = randomTree hp{_maxSize = msz-1} True-        go n msz (Fix (Uni f t)) = Fix . Uni f <$> go (n-1) (msz-1) t-        go n msz (Fix (Bin op l r)) = do-          let nl = countNodes l-              nr = countNodes r-          if nl <= n - 1-             then Fix . Bin op l <$> go (n-nl-1) (msz-nl-1) r-             else do l' <- go (n-1) (msz-nr-1) l-                     pure $ Fix $ Bin op l' r--crossover :: HyperParams -> Individual -> Individual -> Rng (Maybe Individual)-crossover hp ind1 ind2 = do-  b <- state random-  if b < (_pc hp)-     then do let n1 = countNodes $ _tree ind1-                 n2 = countNodes $ _tree ind2-             p1 <- state $ randomR (0, n1-1)-             p2 <- state $ randomR (0, n2-1)-             let part1 = pickLeft p1 $ _tree ind1-                 part2 = pickRight p2 $ _tree ind2-                 t = part1 part2-                 n = countNodes t-             if n <= _maxSize hp-                then pure . Just $ ind1{_tree = t}-                else pure Nothing-     else pure Nothing-  where-    pickRight :: Int -> Fix SRTree -> Fix SRTree-    pickRight 0 node = node-    pickRight n (Fix (Uni f t)) = pickRight (n-1) t-    pickRight n (Fix (Bin op l r)) = let nl = countNodes l-                                     in if nl <= n-1-                                           then pickRight (n-nl-1) r-                                           else pickRight (n-1) l-    pickLeft :: Int -> Fix SRTree -> (Fix SRTree -> Fix SRTree)-    pickLeft 0 node = \t -> t-    pickLeft n (Fix (Uni f t)) = let g = pickLeft (n-1) t in \t' -> Fix $ Uni f (g t')-    pickLeft n (Fix (Bin op l r)) = let nl = countNodes l-                                    in if nl <= n-1-                                          then let g = pickLeft (n-nl-1) r in \t -> Fix $ Bin op l (g t)-                                          else let g = pickLeft (n-1) l in \t -> Fix $ Bin op (g t) r---evolve :: HyperParams -> FitFun -> V.Vector Individual -> Rng Individual-evolve hp fitFun pop = do -    parent1 <- tournament hp pop-    parent2 <- tournament hp pop -    mChild <- crossover hp parent1 parent2-    child' <- case mChild of-                Nothing    -> mutate hp parent1-                Just child -> mutate hp child-    --let p = countParams (_tree child')-    --theta' <- M.fromList compMode <$> replicateM p (randomRange (-1,1))-    case child' of-        Nothing -> pure parent1-        Just c  -> fitFun c--printFinal dist ind dataTrains dataTests = do-  let tree     = relabelParams $ _tree ind-      thetas   = _params ind-      mseTrain = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTrains-      mseTest  = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTests-      r2Train  = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTrains-      r2Test   = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTests-      thetaStr = intercalate "_" $ map (intercalate ";" . map show . M.toList) thetas-  putStrLn "id,Expression,theta,size,MSE_train,MSE_test,R2_train,R2_test"-  putStr $ "0," <> showExpr tree <> "," <> thetaStr <> "," <> show (countNodes tree) <> "," <> show mseTrain <> "," <> show mseTest <> "," <> show r2Train <> "," <> show r2Test--report :: Int -> V.Vector Individual -> IO ()-report gen = mapM_ reportOne-  where reportOne ind = do putStr (show gen)-                           putStr ": "-                           putStr (showExpr (_tree ind))-                           putStr " - " -                           putStr (show (_fit ind))-                           putStr " "-                           print (map M.toList $ _params ind)-{-# INLINE report #-}--evolution :: Int -> HyperParams -> FitFun -> Rng (Individual)-evolution gen hp fitFun = do -    pop <- initialPop hp fitFun-    --liftIO $ report (-1) pop-    go gen pop-        where -            go 0 !pop = pure $ pop  V.! 0-            go n !pop = do-                let best = V.maximumOn _fit $ V.filter (not.isNaN._fit) pop-                pop' <- V.modify (\v -> MV.write v 0 best) <$> V.replicateM (_popSize hp) (evolve hp fitFun pop)-                --liftIO $ report (gen-n) pop'-                go (n-1) pop'

@@ -1,50 +0,0 @@-module Initialization where--data InitiMethod = GROW | FULL | BTC | HALFHALF-{--btc = undefined --def btc(pset_, depth_, length_, type_=None):-    if type_ is None:-        type_ = pset_.ret--    expr = []--    arities = list(map(lambda x: x.arity, pset_.primitives[type_]))-    minFunctionArity = min(arities)-    maxFunctionArity = max(arities)--    # adapt length to restrictions of the primitive set-    if length_ % 2 == 0 and minFunctionArity > 1:-        length_ = length_ + 1 if np.random.random_sample(1) > 0.5 else length_ - 1--    targetLength = length_ - 1 # don't count the root node -    maxFunctionArity = min(maxFunctionArity, targetLength)-    minFunctionArity = min(minFunctionArity, targetLength)-    root = sampleChild(pset_, minFunctionArity, maxFunctionArity, type_) --    # inner lists of the form [node, depth, childIndex] -    # childIndex is only used at the end to transform -    # the representation from breadth to prefix-    expr.append([root, 0, 1])--    openSlots = root.arity --    for i in range(0, length_):-        (node, nodeDepth, childIndex) = expr[i]-        childDepth = nodeDepth + 1-        -        for j in range(0, getArity(node)):-            maxArity = 0 if childDepth == depth_ - 1 else min(maxFunctionArity, targetLength - openSlots)-            minArity = min(minFunctionArity, maxArity)-            child = sampleChild(pset_, minArity, maxArity, type_)--            if j == 0:-                expr[i][2] = len(expr)--            expr.append([child, childDepth, 0])-            openSlots += getArity(child) --    nodes = breadthToPrefix(expr)-    return nodes-    -}

@@ -1,116 +0,0 @@-module Main (main) where--import GP ( HyperParams(HP), fitnessMV, evolution, printFinal )-import Data.SRTree-import System.Random ( getStdGen )-import Control.Monad.State.Strict ( evalStateT )-import Data.SRTree.Datasets ( loadDataset ) -import Options.Applicative-import Data.Massiv.Array -import Util-import Algorithm.SRTree.Likelihoods-import Data.SRTree.Datasets---- Data type to store command line arguments-data Args = Args-  { dataset   :: String,-    _testData :: String,-    popSize   :: Int,-    gens      :: Int,-    _maxSize  :: Int,-    pc        :: Double,-    pm        :: Double,-    _nonterminals :: String,-    _nTournament  :: Int,-    _distribution :: Distribution-  }-  deriving (Show)----- parser of command line arguments-opt :: Parser Args-opt = Args-   <$> strOption-       ( long "dataset"-       <> short 'd'-       <> metavar "INPUT-FILE"-       <> help "CSV dataset." )-   <*> strOption-       ( long "test"-       <> value ""-       <> metavar "INPUT-FILE"-       <> help "CSV dataset." )-   <*> option auto-       ( long "population"-       <> short 'p'-       <> metavar "POP-SIZE"-       <> showDefault-       <> value 100-       <> help "Population size." )-   <*> option auto-      ( long "generations"-      <> short 'g'-      <> metavar "GENS"-      <> showDefault-      <> value 100-      <> help "Number of generations." )-   <*> option auto-      ( long "max-size"-      <> metavar "SIZE"-      <> showDefault-      <> value 20-      <> help "maximum expression size." )-   <*> option auto-      ( long "probCx"-      <> metavar "PC"-      <> showDefault-      <> value 0.9-      <> help "Crossover probability." )-   <*> option auto-      ( long "probMut"-      <> metavar "PM"-      <> showDefault-      <> value 0.3-      <> help "Mutation probability." )-   <*> strOption-       ( long "non-terminals"-       <> value "Add,Sub,Mul,Div,PowerAbs,Recip"-       <> showDefault-       <> help "set of non-terminals to use in the search."-       )-   <*> option auto-       ( long "tournament-size"-       <> value 2-       <> showDefault-       <> help "tournament size."-       )-   <*> option auto-       ( long "distribution"-       <> value MSE-       <> showDefault-       <> help "distribution of the data.")--nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> Fix $ Bin PowerAbs l r), Left recip, Left log, Left exp, Left (\t -> Fix $ Uni SqrtAbs t)]--main :: IO ()-main = do-  args <- execParser opts-  g <- getStdGen-  --(x, y, _) <- loadTrainingOnly (dataset args) True-  --(x_test, y_test, _) <- loadTrainingOnly (_testData args) True--  let datasets = words (dataset args)-  dataTrains <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets-  dataTests  <- if null (_testData args)-                  then pure dataTrains-                  else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args)--  let hp = HP 3 10 (_maxSize args) (popSize args) (_nTournament args) (pc args) (pm args) terms (parseNonTerms $ _nonterminals args)-      (Sz2 _ nFeats) = size . getX $ head dataTrains-      terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]]-  best <- evalStateT (evolution (gens args) hp (fitnessMV (_distribution args) dataTrains)) g-  printFinal (_distribution args) best dataTrains dataTests-  where-    opts = info (opt <**> helper)-            ( fullDesc <> progDesc "Very simple example of GP using SRTree."-           <> header "tinyGP - a very simple example of GP using SRTRee." )

@@ -1,63 +0,0 @@-{-# LANGUAGE  BlockArguments #-}-{-# LANGUAGE  TupleSections #-}--module Util where--import qualified Data.Map.Strict as Map-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.SRTree-import Data.SRTree.Eval-import Algorithm.SRTree.Opt-import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Build-import Algorithm.EqSat.Info--import Algorithm.SRTree.NonlinearOpt-import System.Random-import Algorithm.SRTree.Likelihoods---import Algorithm.SRTree.ModelSelection---import Algorithm.SRTree.Opt-import qualified Data.IntMap.Strict as IM-import Control.Monad.State.Strict-import Control.Monad ( when, replicateM, forM, forM_ )-import Data.Maybe ( fromJust )-import Data.List ( maximumBy )-import Data.Function ( on )-import List.Shuffle ( shuffle )-import Data.List.Split ( splitOn )-import Data.Char ( toLower )-import qualified Data.IntSet as IntSet-import Data.SRTree.Datasets-import Algorithm.EqSat.Queries---    {--type DataSet = (SRMatrix, PVector, Maybe PVector)---getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)-getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)--getX :: DataSet -> SRMatrix-getX (a, _, _) = a--getTarget :: DataSet -> PVector-getTarget (_, b, _) = b--getError :: DataSet -> Maybe PVector-getError (_, _, c) = c--loadTrainingOnly fname b = getTrain <$> loadDataset fname b--}--parseNonTerms = Prelude.map toNonTerm . splitOn ","-  where-    binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]-    uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]-    toNonTerm xs' = let xs = Prelude.map toLower xs'-                    in case binTerms Map.!? xs of-                          Just op -> Right $ \l r -> Fix $ Bin op l r-                          Nothing -> case uniTerms Map.!? xs of-                                          Just f -> Left $ \t -> Fix $ Uni f t-                                          Nothing -> error $ "invalid non-terminal " <> show xs-

@@ -24,150 +24,344 @@ import Data.Function (on) import Data.IntMap (IntMap) import qualified Data.IntMap as IntMap-import Data.List (intercalate, minimumBy)+import qualified Data.IntSet as IntSet+import Data.List (intercalate) import Data.Map (Map) import qualified Data.Map as Map import Data.Maybe (mapMaybe) import Data.SRTree import Data.HashSet (HashSet) import qualified Data.HashSet as Set-import Control.Monad ( zipWithM )--import Debug.Trace+import Control.Monad ( zipWithM, forM_ )  -- | The `Scheduler` stores a map with the banned iterations of a certain rule .  -- TODO: make it more customizable. type Scheduler a = State (IntMap Int) a  -- to avoid importing-fromJust :: Maybe a -> a-fromJust (Just x) = x-fromJust _        = error "fromJust called with Nothing"-{-# INLINE fromJust #-}- -- | runs equality saturation from an expression tree, -- a given set of rules, and a cost function. -- Returns the tree with the smallest cost.-eqSat :: Monad m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree)+eqSat :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) eqSat expr rules costFun maxIt =     do root <- fromTree costFun expr-       (end, it) <- runEqSat costFun rules maxIt-       best      <- getBestExpr root-       --info      <- gets ((IntMap.! root) . _eClass)-       --info2     <- gets ((IntMap.! 9) . _eClass)-       --traceShow (info, info2) $-       if not end -- if had an early stop-         then do modify' (const emptyGraph) >> eqSat best rules costFun it -- reapplies eqsat on the best so far-         else pure best+       _ <- runEqSat costFun rules maxIt+       recalculateBest costFun root  type CostMap = Map EClassId (Int, Fix SRTree)  -- | recalculates the costs with a new cost function-recalculateBest :: Monad m => CostFun -> EClassId -> EGraphST m (Fix SRTree)+recalculateBest :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree) recalculateBest costFun eid =-    do classes <- gets _eClass-       let costs = fillUpCosts classes Map.empty+    do ecls <- allClasses+       let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls]+           costs   = fillUpCosts classes Map.empty        eid' <- canonical eid-       pure $ snd $ costs Map.! eid'+       case Map.lookup eid' costs of+         Just (_, t) -> pure t+         Nothing     -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid'+                              <> " nClasses=" <> show (IntMap.size classes)+                              <> " costSize=" <> show (Map.size costs)     where-        nodeCost :: CostMap -> ENode -> Maybe (Int, Fix SRTree)+        nodeCost :: CostMap -> ENode -> (Int, Fix SRTree)         nodeCost costMap enode =-          do optChildren <- traverse (costMap Map.!?) (childrenOf enode) -- | gets the cost of the children, if one is missing, returns Nothing-             let cc = map fst optChildren-                 nc = map snd optChildren-                 n  = replaceChildren cc enode-                 c  = costFun n-             pure (c + sum cc, Fix $ replaceChildren nc enode) -- | otherwise, returns the cost of the node + children and the expression so far--        minimumBy' f [] = Nothing-        minimumBy' f xs = Just $ minimumBy f xs+          -- A child that has not been costed yet (a cycle, or a class whose+          -- cost is computed later in this iteration) contributes a large+          -- sentinel instead of 0: a 0 placeholder is cheaper than the real+          -- cost, so the fixpoint below would keep the stale placeholder tree+          -- (e.g. `x * 0.0` for `x * (y + z)`). Real costs always beat it.+          let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (costMap Map.!? cid) | cid <- eChildren enode ]+              c  = case enode of+                     ENAry op _ -> costFun (Bin (toOp op) 0 0)+                     _          -> costFun (replaceChildren cc (fromENode enode))+          in (c + sum cc, Fix $ case enode of+                 ENAry op _ -> unfix (naryTree op nc)+                 _          -> replaceChildren nc (fromENode enode)) -- | missing children (cyclic classes) get cost 0 so every class is costed+        costSentinel :: Int+        costSentinel = 1000000          fillUpCosts :: IntMap EClass -> CostMap -> CostMap-        fillUpCosts classes m =-            case IntMap.foldrWithKey costOfClass (False, m) classes of -- applies costOfClass to each class-              (False, _) -> m-              (True, m') -> fillUpCosts classes m' -- | if something changed, recurse+        fillUpCosts classes = go (IntMap.size classes + 1) (IntMap.keysSet classes)+          where+            go 0 _ m = m+            go n dirty m+              | IntSet.null dirty = m+              | otherwise = go (n - 1) dirty' m'+              where+                (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty+                step (d, cm) eid = case IntMap.lookup eid classes of+                  Nothing -> (d, cm)+                  Just ecl ->+                    let currentCost = Map.lookup eid cm+                        minCost     = Set.foldl' (\acc en -> let c = nodeCost cm en+                                                  in case acc of+                                                    Nothing  -> Just c+                                                    Just c'  -> Just (if fst c <= fst c' then c else c')+                                                ) Nothing (_eNodes ecl)+                        (changed, cm') = case (currentCost, minCost) of+                          (_, Nothing)            -> (False, cm)+                          (Nothing, Just new)     -> (True, Map.insert eid new cm)+                          (Just old, Just new)+                            | fst old <= fst new  -> (False, cm)+                            | otherwise           -> (True, Map.insert eid new cm)+                        d' = if changed+                             then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl)+                             else d+                    in d' `seq` cm' `seq` (d', cm') -        costOfClass :: EClassId -> EClass -> (Bool, CostMap) -> (Bool, CostMap)-        costOfClass eid ecl (b, m) =-            let currentCost = m Map.!? eid-                minCost     = minimumBy' (compare `on` fst)  -- get the minimum available cost of the nodes of this class-                            $ mapMaybe (nodeCost m)-                            $ map decodeEnode-                            $ Set.toList (_eNodes ecl)-            in case (currentCost, minCost) of -- replace the costs accordingly-                  (_, Nothing)         -> (b, m)-                  (Nothing, Just new)  -> (True, Map.insert eid new m)-                  (Just old, Just new) -> if fst old <= fst new-                                            then (b, m)-                                            else (True, Map.insert eid new m)+-- | Recompute every e-class's cost-minimal @_best@/_cost@ bottom-up and write+-- it back into the graph. Needed after loading a graph whose best/cost were+-- not persisted (e.g. via srtree-db), where @_best@ may otherwise hold an+-- arbitrary (potentially large) e-node.+recalculateBestAll :: ClassStore m => CostFun -> EGraphST m ()+recalculateBestAll costFun = do+  ecls <- allClasses+  let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls]+      bests = fixpoint classes IntMap.empty+  forM_ (IntMap.toList bests) $ \(eid, (c, en)) ->+    case IntMap.lookup eid classes of+      Nothing -> pure ()+      Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } }+  where+    nodeCost :: IntMap (Int, ENode) -> ENode -> (Int, ENode)+    nodeCost cm en =+      let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ]+          c  = case en of+                 ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc+                 _          -> costFun (replaceChildren cc (fromENode en)) + sum cc+      in (c, en)+    costSentinel :: Int+    costSentinel = 1000000 +    fixpoint :: IntMap EClass -> IntMap (Int, ENode) -> IntMap (Int, ENode)+    fixpoint classes0 = go (IntMap.size classes0 + 1) (IntMap.keysSet classes0)+      where+        go 0 _ m        = m+        go n dirty m+          | IntSet.null dirty = m+          | otherwise = go (n - 1) dirty' m'+          where+            (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty+            step (d, cm) eid = case IntMap.lookup eid classes0 of+              Nothing -> (d, cm)+              Just ecl ->+                let current = IntMap.lookup eid cm+                    minNode = Set.foldl' (\acc en -> let c = nodeCost cm en+                                                     in case acc of+                                                          Nothing  -> Just c+                                                          Just c'  -> Just (if fst c <= fst c' then c else c'))+                                         Nothing (_eNodes ecl)+                    (changed, cm') = case (current, minNode) of+                      (_, Nothing)        -> (False, cm)+                      (Nothing, Just new) -> (True, IntMap.insert eid new cm)+                      (Just old, Just new)+                        | fst old <= fst new -> (False, cm)+                        | otherwise          -> (True, IntMap.insert eid new cm)+                    d' = if changed+                         then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl)+                          else d+                 in d' `seq` cm' `seq` (d', cm')++-- | Like 'recalculateBestAll' but streamed: each e-class body is fetched on+-- demand through 'ClassStore' (so a paged graph never materializes every body+-- at once) and only the small @(cost, best e-node)@ map is kept resident. The+-- structural worklist fixpoint is identical.+recalculateBestAllStream :: ClassStore m => CostFun -> EGraphST m ()+recalculateBestAllStream costFun = do+  ids <- allKeys+  let idSet = IntSet.fromList ids+      costSentinel = 1000000+      nodeCost cm en =+        let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ]+            c  = case en of+                   ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc+                   _          -> costFun (replaceChildren cc (fromENode en)) + sum cc+        in (c, en)+      stepEid cm eid = do+        mec <- readDirect eid+        case mec of+          Nothing -> pure (IntSet.empty, cm)+          Just ecl -> do+            let current = IntMap.lookup eid cm+                minNode = Set.foldl' (\acc en -> let c = nodeCost cm en+                                                 in case acc of+                                                      Nothing  -> Just c+                                                      Just c'  -> Just (if fst c <= fst c' then c else c'))+                                    Nothing (_eNodes ecl)+                (changed, cm') = case (current, minNode) of+                  (_, Nothing)        -> (False, cm)+                  (Nothing, Just new) -> (True, IntMap.insert eid new cm)+                  (Just old, Just new)+                    | fst old <= fst new -> (False, cm)+                    | otherwise          -> (True, IntMap.insert eid new cm)+                dirty = if changed+                          then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) IntSet.empty (_parents ecl)+                          else IntSet.empty+            pure (dirty, cm')+      fixpoint n dirty cm+        | n <= 0 || IntSet.null dirty = pure cm+        | otherwise = go (IntSet.toList dirty) IntSet.empty cm+        where+          go [] d acc = fixpoint (n - 1) d acc+          go (e : es) d acc = do+            (d', m') <- stepEid acc e+            go es (IntSet.union d d') m'+  cm <- fixpoint (IntSet.size idSet + 1) idSet IntMap.empty+  forM_ (IntMap.toList cm) $ \(eid, (c, en)) -> do+    mec <- readDirect eid+    case mec of+      Nothing -> pure ()+      Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } }++-- | Streaming variant of 'recalculateBest': computes the cost-minimal tree for a+-- single root without materializing every e-class body at once.+recalculateBestStream :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree)+recalculateBestStream costFun eid = do+  ids <- allKeys+  let idSet = IntSet.fromList ids+      costSentinel = 1000000+      nodeCost cm en =+        let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (Map.lookup cid cm) | cid <- eChildren en ]+            c  = case en of+                   ENAry op _ -> costFun (Bin (toOp op) 0 0)+                   _          -> costFun (replaceChildren cc (fromENode en))+        in (c + sum cc, Fix $ case en of+               ENAry op _ -> unfix (naryTree op nc)+               _          -> replaceChildren nc (fromENode en))+      stepEid cm eid' = do+        mec <- lookupClass eid'+        case mec of+          Nothing -> pure (IntSet.empty, cm)+          Just ecl -> do+            let current = Map.lookup eid' cm+                minCost = Set.foldl' (\acc en -> let c = nodeCost cm en+                                                 in case acc of+                                                      Nothing -> Just c+                                                      Just c' -> Just (if fst c <= fst c' then c else c'))+                                   Nothing (_eNodes ecl)+                (changed, cm') = case (current, minCost) of+                  (_, Nothing) -> (False, cm)+                  (Nothing, Just new) -> (True, Map.insert eid' new cm)+                  (Just old, Just new)+                    | fst old <= fst new -> (False, cm)+                    | otherwise -> (True, Map.insert eid' new cm)+                dirty = if changed+                          then Set.foldl' (\acc (pid,_) -> IntSet.insert pid acc) IntSet.empty (_parents ecl)+                          else IntSet.empty+            pure (dirty, cm')+      fixpoint n dirty cm+        | n <= 0 || IntSet.null dirty = pure cm+        | otherwise = go (IntSet.toList dirty) IntSet.empty cm+        where+          go [] d acc = fixpoint (n - 1) d acc+          go (e : es) d acc = do+            (d', m') <- stepEid acc e+            go es (IntSet.union d d') m'+  cm <- fixpoint (IntSet.size idSet + 1) idSet Map.empty+  eid' <- canonical eid+  case Map.lookup eid' cm of+    Just (_, t) -> pure t+    Nothing -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid'+                     <> " costSize=" <> show (Map.size cm)++-- | Run equality saturation and stream the final extraction (see+-- 'recalculateBestStream'), so a paged graph is never fully materialized.+eqSatStream :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree)+eqSatStream expr rules costFun maxIt = do+  root <- fromTree costFun expr+  _ <- runEqSat costFun rules maxIt+  recalculateBestAllStream costFun+  recalculateBestStream costFun root++-- | replaces the equality rules with two one-way rules+replaceEqRules :: Rule -> [Rule]+replaceEqRules (p1 :=> p2)  = [p1 :=> p2]+replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1]+replaceEqRules (r :| cond)  = map (:| cond) $ replaceEqRules r++-- | Compile a rule source into a query, or `Nothing` for n-ary patterns that+-- use the direct multiset matcher instead.+compileSource :: Rule -> Maybe (Query, [ClassOrVar], ClassOrVar)+compileSource r = if hasNAry (source r)+                    then Nothing+                    else Just (compileToQuery (source r))++-- | Cap on the total number of rule matches applied in a single eqsat+-- iteration. Combined with the per-rule caps ('ruleBudget'/'ruleRootVisit' for+-- n-ary, 'ruleMatchBudget' for the cached path) and the persistent+-- mark-on-attempt seen-set (which makes each rule's budget advance to new+-- matches), this bounds a single iteration's apply/rebuild work regardless of+-- graph size.+iterMatchBudget :: Int+iterMatchBudget = 2000+ -- | run equality saturation for a number of iterations-runEqSat :: Monad m => CostFun -> [Rule] -> Int -> EGraphST m (Bool, Int)-runEqSat costFun rules maxIter = go maxIter IntMap.empty+runEqSat :: ClassStore m => CostFun -> [Rule] -> Int -> EGraphST m (Bool, Int)+runEqSat costFun rules maxIter = go maxIter IntMap.empty compiledRules     where         rules' = concatMap replaceEqRules rules--        -- replaces the equality rules with two one-way rules-        replaceEqRules :: Rule -> [Rule]-        replaceEqRules (p1 :=> p2)  = [p1 :=> p2]-        replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1]-        replaceEqRules (r :| cond)  = map (:| cond) $ replaceEqRules r+        compiledRules = map (\r -> (r, compileSource r)) rules' -        go it sch = do eNodes   <- gets _eNodeToEClass-                       eClasses <- gets _eClass-                       --createDB -- TODO: partial db is still incomplete -                       --db       <- gets (_patDB . _eDB) -- createDB -- creates the DB+        go it sch compiled =+          do -- reset dirty flag before processing this iteration+             modify' $ over (eDB . changed) (const False) -                       -- step 1: match the rules-                       let matchSch        = matchWithScheduler it-                           matchAll        = zipWithM matchSch [0..]-                           (rules, sch')   = runState (matchAll rules') sch+             -- step 1: match the rules using cached compiled queries+             let matchSch  = matchWithScheduler it+                 adapted i (r, cq) = map (,cq) <$> matchSch i r+                 matchAll  = zipWithM adapted [0..]+                 (filtered, sch') = runState (matchAll compiled) sch -                       -- step 2: apply matches and rebuild-                       matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rules-                       mapM_ (uncurry (applyMatch costFun)) $ concat matches-                       rebuild costFun+             -- step 2: apply matches and rebuild+             matches <- mapM (\(rule, cq) -> map (rule,) <$> case cq of+                                Just q  -> do paged <- isPagedGraph+                                              if paged+                                                then matchStreamCached (Just (show (source rule))) (source rule)+                                                else matchCachedWith (Just (show (source rule))) q+                                Nothing -> matchSaturated (source rule)) $ concat filtered+             -- bound the total number of matches applied per iteration so a+             -- single iteration's apply/rebuild work stays bounded on huge+             -- graphs (genuine matches; we just process them over more iters).+             mapM_ (uncurry (applyMatch costFun)) (take iterMatchBudget (concat matches))+             rebuild costFun -                       -- recalculate heights-                       --calculateHeights-                       eNodes'   <- gets _eNodeToEClass-                       eClasses' <- gets _eClass+             -- check dirty flag: if no modifications occurred, we've saturated+             changed <- gets (_changed . _eDB)+             if it == 1 || not changed+                then pure (True, it)+                 else+                   do eClasses <- gets _eClass+                      if IntMap.size eClasses > 1500+                        then throttle it sch' compiled+                        else go (it-1) sch' compiled -                       -- if nothing changed, return-                       if it == 1 || (eNodes' == eNodes && eClasses' == eClasses)-                          then pure (True, it)-                          else if IntMap.size eClasses' > 1500 -- maximum allowed number of e-classes. TODO: customize-                                 then pure (False, it)-                                 else go (it-1) sch'+        throttle it sch compiled = do+          cleanMaps+          eClasses <- gets _eClass+          if IntMap.size eClasses <= 1500+            then go (it-1) sch compiled+            else do applySingleMergeOnlyEqSat costFun rules+                    changed <- gets (_changed . _eDB)+                    if it <= 1 || not changed+                      then pure (False, it)  -- give up and return early stop+                      else throttle (it-1) sch compiled  -- | apply a single step of merge-only equality saturation-applySingleMergeOnlyEqSat :: Monad m => CostFun -> [Rule] -> EGraphST m ()+applySingleMergeOnlyEqSat :: ClassStore m => CostFun -> [Rule] -> EGraphST m () applySingleMergeOnlyEqSat costFun rules =-  do db <- gets (_patDB . _eDB) -- createDB-     let matchSch        = matchWithScheduler 10+  do let matchSch        = matchWithScheduler 10          matchAll        = zipWithM matchSch [0..]-         (rls, sch')     = runState (matchAll rules') IntMap.empty-     --matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rls-     --mapM_ (uncurry (applyMergeOnlyMatch costFun)) $ take 500 $ concat matches+         (rls, _)        = runState (matchAll rules') IntMap.empty      matches <- getNMatches 500 rls      rebuild costFun-     -- recalculate heights-     --calculateHeights       where         rules' = concatMap replaceEqRules rules -        -- replaces the equality rules with two one-way rules-        replaceEqRules :: Rule -> [Rule]-        replaceEqRules (p1 :=> p2)  = [p1 :=> p2]-        replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1]-        replaceEqRules (r :| cond)  = map (:| cond) $ replaceEqRules r-         getNMatches n []       = pure []         getNMatches 0 _        = pure []         getNMatches n ([]:rss) = getNMatches n rss-        getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> match (source r)-                                        let (x, y) = splitAt n matches+        getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> matchSaturated (source r)+                                        let (x, _) = splitAt n matches                                             m      = length x                                         if m == n                                            then pure matches@@ -179,7 +373,7 @@ matchWithScheduler :: Int -> Int -> Rule -> Scheduler [Rule] -- [(Rule, (Map ClassOrVar ClassOrVar, ClassOrVar))] matchWithScheduler it ruleNumber rule =   do mbBan <- gets (IntMap.!? ruleNumber)-     if mbBan /= Nothing && fromJust mbBan <= it -- check if the rule is banned+     if maybe False (<= it) mbBan -- check if the rule is banned         then pure []         else do -- let matches = match db (source rule)                 modify (IntMap.insert ruleNumber (it+5))

@@ -20,76 +20,161 @@ import System.Random (Random (randomR), StdGen) import Control.Lens ( over ) import Control.Monad ( forM_, when, foldM, forM )-import Data.Maybe ( fromMaybe, catMaybes )+import Data.Maybe import Data.SRTree import Algorithm.EqSat.Egraph---import Algorithm.EqSat.Info import Algorithm.EqSat.DB import qualified Data.IntMap.Strict as IntMap+import Data.IntMap.Strict (IntMap) import Data.Map.Strict ( Map ) import qualified Data.Map.Strict as Map+import qualified Data.HashMap.Strict as HashMap import qualified Data.HashSet as Set import Control.Monad.State.Strict import Control.Monad.Identity+import GHC.Stack (HasCallStack)+ import Data.SRTree.Recursion (cataM)+import Data.List (sort) import Algorithm.EqSat.Info import qualified Data.IntSet as IntSet-import Data.Maybe-import Data.Sequence (Seq(..), (><))-import Data.List ( nub )-import Debug.Trace (trace, traceShow) +import qualified Data.Set as RangeSet++ -- | adds a new or existing e-node (merging if necessary)-add :: Monad m => CostFun -> ENode -> EGraphST m EClassId-add costFun enode =-  do enode''   <- canonize enode                                             -- canonize e-node-     constEnode <- calculateConsts enode''-     enode' <- case constEnode of-                 ConstVal x -> pure $ Const x-                 ParamIx  x -> pure $ Param x-                 _          -> case enode'' of-                                 Bin Sub c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)-                                                     pure $ case constType of-                                                              ParamIx x -> Bin Add c1 c2-                                                              _         -> enode''-                                 Bin Div c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)-                                                     pure $ case constType of-                                                              ParamIx x -> Bin Mul c1 c2-                                                              _         -> enode''-                                 _             -> pure $ enode''+add :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m EClassId+add costFun enode = do+  enode''  <- canonize enode+  enode''' <- foldConsts costFun enode'' -     maybeEid <- gets ((Map.!? enode') . _eNodeToEClass)                -- check if canonical e-node exists-     case maybeEid of+  maybeEid <- lookupNode enode'''+  case maybeEid of        Just eid -> pure eid        Nothing  -> do          curId <- gets (_nextId . _eDB)                             -- get the next available e-class id-         modify' $ over canonicalMap (IntMap.insert curId curId)           -- insert e-class id into canon map-                 . over eNodeToEClass (Map.insert enode' curId)     -- associate new e-node with id-                 . over (eDB . nextId) (+1)                                -- update next id-                 . over (eDB . worklist) (Set.insert (curId, enode'))      -- add e-node and id into worklist-         forM_ (childrenOf enode') (addParents curId enode')        -- update the children's parent list-         info <- makeAnalysis costFun enode'-         h    <- getChildrenMinHeight enode'-         let newClass = createEClass curId enode' info h              -- create e-class-         modify' $ over eClass (IntMap.insert curId newClass)              -- insert new e-class into e-graph+         insertCanonical curId curId                                 -- register the class as its own representative+         insertNode enode''' curId                                  -- associate new e-node with id (bounded on paged graphs)+         modify' $ over (eDB . nextId) (+1)                          -- update next id+                 . over (eDB . worklist) (Set.insert (curId, enode'''))      -- add e-node and id into worklist+         forM_ (eChildren enode''') (addParents curId enode''')        -- update the children's parent list+         info <- makeAnalysis costFun enode'''+         h    <- getChildrenMinHeight enode'''+         let newClass = createEClass curId enode''' info h              -- create e-class+         -- insert via 'insertClass' so a paged (DB-backed) class store also+         -- persists the new class's page; for a pure graph this is identical+         -- to inserting into @_eClass@ directly.+         insertClass newClass          --modifyEClass costFun curId                                 -- simplify eclass if it evaluates to a number           -- update database-         addToDB enode' curId                                       -- add new node to db-         modify' $ over (eDB . sizeDB)-                 $ IntMap.insertWith (IntSet.union) (_size info) (IntSet.singleton curId)+         addToDB enode''' curId                                       -- add new node to db+         tracking <- gets (_trackDBs . _eDB)+         when tracking $+           modify' $ over (eDB . sizeDB)+                   $ IntMap.insertWith (IntSet.union) (_size info) (IntSet.singleton curId)          modify' $ over (eDB . unevaluated) (IntSet.insert curId)+                 . over (eDB . changed) (const True)          pure curId   where-    addParents :: Monad m => EClassId -> ENode -> EClassId -> EGraphST m ()+    addParents :: ClassStore m => EClassId -> ENode -> EClassId -> EGraphST m ()     addParents cId node c =       do ec <- getEClass c          let ec' = ec{ _parents = Set.insert (cId, node) (_parents ec) }-         modify' $ over eClass (IntMap.insert c ec')+         -- write through 'insertClass' so a paged store keeps the updated parents+         insertClass ec' +-- | Add a binary (SRTree-based) node, converting it to a flattened ENode.+-- Sub and Div are canonicalized away at insertion: `x - y` becomes+-- `x + (-1)*y` and `x / y` becomes `x * recip y`, so no Sub/Div e-node ever+-- enters the e-graph and the Sub/Div-aware rules become redundant.+addTree :: (ClassStore m, HasCallStack) => CostFun -> SRTree EClassId -> EGraphST m EClassId+addTree costFun (Bin Sub l r) = do+  neg <- addNegate costFun r+  add costFun =<< mkENary EAdd [l, neg]+addTree costFun (Bin Div l r) = do+  rec <- add costFun (EUni Recip r)+  add costFun =<< mkENary EMul [l, rec]+addTree costFun t = toENode t >>= add costFun+{-# INLINE addTree #-}++-- | builds the e-class for the negation of the e-class `t`, represented as+-- `(-1) * t` (matching the pattern-level `negate` encoding in Algorithm.EqSat.DB).+addNegate :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId+addNegate costFun t = do+  negOne <- add costFun (EConst (-1))+  add costFun =<< mkENary EMul [negOne, t]++-- | Fused 'calculateConsts' + 'foldConstants': fetches each child's constant+-- info a single time, detects fully-constant nodes (replaced by EConst/EParam)+-- and folds together all-but-one constant children of an ENAry+-- (e.g. 2+3+x becomes 5+x). Constants that are already folded single subtrees+-- are handled by the same child-constant walk.+foldConsts :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode+foldConsts _ en@(ENAry _ m) | IntMap.null m = pure en+foldConsts costFun en@(ENAry op m) = do+  let xs = expandedList m+  infos <- mapM (fmap (_consts . _info) . getEClass) xs+  case foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) infos of+    ConstVal x -> pure (EConst x)+    ParamIx x  -> pure (EParam x)+    _          -> foldENary costFun op m infos+foldConsts _ en = do+  infos <- mapM (fmap (_consts . _info) . getEClass) (eChildren en)+  case combineConsts (replaceChildren infos (fromENode en)) of+    ConstVal x -> pure (EConst x)+    ParamIx x  -> pure (EParam x)+    _          -> pure en+{-# INLINE foldConsts #-}++-- | Fold together all-but-one constant children of an ENAry multiset.+foldENary :: (ClassStore m, HasCallStack) => CostFun -> NOp -> IntMap Int -> [Consts] -> EGraphST m ENode+foldENary costFun op m infos = do+  let xs = expandedList m+      (consts, rest) = foldr step ([], []) (zip xs infos)+      step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs)+      step (x, _)          (cs, rs)              = (cs, x:rs)+  if length consts >= 2+    then do+      let folded = case op of+                     EAdd -> sum consts+                     EMul -> product consts+      if isNaN folded || isInfinite folded+        then pure (ENAry op m)+        else do+          cid <- add costFun (EConst folded)+          pure (ENAry op (imFromList (cid : rest)))+    else pure (ENAry op m)+{-# INLINE foldENary #-}++-- | Fold together all-but-one constant children of an ENAry at insertion+-- time (e.g. 2+3+x becomes 5+x). Constants that are already folded+-- single subtrees are handled by 'calculateConsts' above; this handles the+-- flattened case where several constant terms land in one multiset.+foldConstants :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode+foldConstants _ en@(ENAry _ m) | IntMap.size m < 2 = pure en+foldConstants costFun en@(ENAry op m) = do+  let xs = expandedList m+  infos <- mapM (fmap (_consts . _info) . getEClass) xs+  let (consts, rest) = foldr step ([], []) (zip xs infos)+      step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs)+      step (x, _)          (cs, rs)              = (cs, x:rs)+  if length consts >= 2+    then do+      let folded = case op of+                     EAdd -> sum consts+                     EMul -> product consts+      if isNaN folded || isInfinite folded+        then pure en+        else do+          cid <- add costFun (EConst folded)+          pure (ENAry op (imFromList (cid : rest)))+    else pure en+foldConstants _ en = pure en+ -- | rebuilds the e-graph after inserting or merging -- e-classes-rebuild :: Monad m => CostFun -> EGraphST m ()+rebuild :: (ClassStore m, HasCallStack) => CostFun -> EGraphST m () rebuild costFun =   do wl <- gets (_worklist . _eDB)      al <- gets (_analysis . _eDB)@@ -102,21 +187,23 @@ -- | repairs e-node by canonizing its children -- if the canonized e-node already exists in -- e-graph, merge the e-classes-repair :: Monad m => CostFun -> EClassId -> ENode -> EGraphST m ()+repair :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> ENode -> EGraphST m () repair costFun ecId enode =-  do modify' $ over eNodeToEClass (Map.delete enode)+  do modify' $ over eNodeToEClass (HashMap.delete enode)      enode'  <- canonize enode      ecId'   <- canonical ecId-     doExist <- gets ((Map.!? enode') . _eNodeToEClass)+     doExist <- lookupNode enode'      case doExist of         Just ecIdCanon -> do mergedId <- merge costFun ecIdCanon ecId'-                             modify' $ over eNodeToEClass (Map.insert enode' mergedId)-        Nothing        -> modify' $ over eNodeToEClass (Map.insert enode' ecId')+                             insertNode enode' mergedId+                             addToDB enode' mergedId+        Nothing        -> do insertNode enode' ecId'+                             addToDB enode' ecId' {-# INLINE repair #-}  -- | repair the analysis of the e-class -- considering the new added e-node-repairAnalysis :: Monad m => CostFun -> EClassId -> ENode -> EGraphST m ()+repairAnalysis :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> ENode -> EGraphST m () repairAnalysis costFun ecId enode =   do ecId'  <- canonical ecId      enode' <- canonize enode@@ -125,15 +212,17 @@      let newData = joinData (_info eclass) info          eclass' = eclass { _info = newData }      when (_info eclass /= newData) $-       do modify' $ over (eDB . analysis) (_parents eclass <>)-                  . over eClass (IntMap.insert ecId' eclass')-                  . over (eDB . refits) (Set.insert ecId')+       do let bestChanged = _best (_info eclass) /= _best newData+          modify' $ over (eDB . analysis) (_parents eclass <>)+                  . (if bestChanged && isJust (_fitness (_info eclass)) then over (eDB . refits) (IntSet.insert ecId') else id)+          -- write through 'insertClass' so a paged store keeps the updated body+          insertClass eclass'           _ <- modifyEClass costFun ecId'           pure () {-# INLINE repairAnalysis #-}  -- | merge to equivalent e-classes-merge :: Monad m => CostFun -> EClassId -> EClassId -> EGraphST m EClassId+merge :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EClassId -> EGraphST m EClassId merge costFun c1 c2 =   do c1' <- canonical c1      c2' <- canonical c2@@ -142,38 +231,43 @@        else do (led, ledC, ledOrig, sub, subC, subOrig) <- getLeaderSub c1' c1 c2' c2  -- the leader will be the e-class with more parents                mergeClasses led ledC ledOrig sub subC subOrig         -- merge sub into leader   where-    mergeClasses :: Monad m => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId+    mergeClasses :: (ClassStore m, HasCallStack) => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId     mergeClasses led ledC ledO sub subC subO =-      do modify' $ over canonicalMap (IntMap.insert sub led . IntMap.insert subO led) -- points sub e-class to leader to maintain consistency-         let -- create new e-class with same id as led-             newC = EClass led-                           (_eNodes ledC `Set.union` _eNodes subC)-                           (_parents ledC <> _parents subC)-                           (min (_height ledC) (_height subC))-                           (joinData (_info ledC) (_info subC))--         modify' $ over eClass (IntMap.insert led newC . IntMap.delete sub) -- delete sub e-class and replace leader-                 . over (eDB . worklist) (_parents subC <>)         -- insert parents of sub into worklist-         when (_info newC /= _info ledC)                            -- if there was change in data,-           $ modify' $ over (eDB . analysis) (_parents ledC <>)     --   insert parents into analysis-                     . over (eDB . refits) (Set.insert led)+      do insertCanonical sub led       -- persist/register the canonical merges+         insertCanonical subO led+         let newC = EClass led+                         (_eNodes ledC `Set.union` _eNodes subC)+                         (_parents ledC <> _parents subC)+                         (min (_height ledC) (_height subC))+                         (joinData (_info ledC) (_info subC))+         forM_ (Set.toList (_eNodes subC)) $ \en -> insertNode en led+         -- write the merged body through the class store (a paged store keeps the+         -- authoritative page) and drop the absorbed class+         insertClass newC+         deleteClass sub+         modify' $ over (eDB . worklist) (_parents subC <>)+         when (_info newC /= _info ledC)+           $ do let bestChanged = _best (_info newC) /= _best (_info ledC)+                modify' $ over (eDB . analysis) (_parents ledC <>)+                           . (if bestChanged && isJust (_fitness (_info ledC)) then over (eDB . refits) (IntSet.insert led) else id)          when (_info newC /= _info subC)            $ modify' $ over (eDB . analysis) (_parents subC <>)-         updateDBs newC led ledC ledO sub subC subO+         tracking <- gets (_trackDBs . _eDB)+         when tracking $ updateDBs newC led ledC ledO sub subC subO          modifyEClass costFun led-         --forM_ (_eNodes newC) $ \en -> addToDB (decodeEnode en) led+         modify' $ over (eDB . changed) (const True)          pure led      getLeaderSub c1 c1O c2 c2O =       do ec1 <- getEClass c1          ec2 <- getEClass c2-         let n1 = length (_parents ec1)-             n2 = length (_parents ec2)+         let n1 = Set.size (_parents ec1)+             n2 = Set.size (_parents ec2)          pure $ if n1 >= n2                   then (c1, ec1, c1O, c2, ec2, c2O)                   else (c2, ec2, c2O, c1, ec1, c1O) -    updateDBs :: Monad m => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m ()+    updateDBs :: (ClassStore m, HasCallStack) => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m ()     updateDBs newC led ledC ledO sub subC subO = do       updateFitnessDB newC led ledC ledO sub subC subO       updateSizeDB newC led ledC ledO sub subC subO@@ -188,20 +282,20 @@      updateFitnessDB :: Monad m => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m ()     updateFitnessDB newC led ledC ledO sub subC subO =-      if (isJust fitNew)-       then do-        when (fitNew /= fitLed) $ do-          if isNothing fitLed-             then modify' $ over (eDB . unevaluated) (IntSet.delete led . IntSet.delete ledO)-             else modify' $ over (eDB . fitRangeDB) (removeRange led (fromJust fitLed) . removeRange ledO (fromJust fitLed))-                          . over (eDB . sizeFitDB) (IntMap.adjust (removeRange ledO (fromJust fitLed) . removeRange led (fromJust fitLed)) szLed)-          modify' $ over (eDB . fitRangeDB) (insertRange led (fromJust fitNew))-                  . over (eDB . sizeFitDB) (IntMap.adjust (insertRange led (fromJust fitNew)) szNew . IntMap.insertWith (><) szNew Empty)-        if isNothing fitSub-           then modify' $ over (eDB . unevaluated) (IntSet.delete sub . IntSet.delete subO)-           else modify' $ over (eDB . fitRangeDB) (removeRange sub (fromJust fitSub) . removeRange subO (fromJust fitSub))-                        . over (eDB . sizeFitDB) (IntMap.adjust (removeRange subO (fromJust fitSub) . removeRange sub (fromJust fitSub)) szSub)-       else modify' $ over (eDB . unevaluated) (IntSet.insert led . IntSet.delete ledO . IntSet.delete sub . IntSet.delete subO)+      case fitNew of+        Nothing -> modify' $ over (eDB . unevaluated) (IntSet.insert led . IntSet.delete ledO . IntSet.delete sub . IntSet.delete subO)+        Just fn -> do+          when (fitNew /= fitLed) $ do+            modify' $ case fitLed of+              Nothing -> over (eDB . unevaluated) (IntSet.delete led . IntSet.delete ledO)+              Just fl -> over (eDB . fitRangeDB) (removeRange led fl . removeRange ledO fl)+                       . over (eDB . sizeFitDB) (IntMap.adjust (removeRange ledO fl . removeRange led fl) szLed)+            modify' $ over (eDB . fitRangeDB) (insertRange led fn)+                    . over (eDB . sizeFitDB) (IntMap.adjust (insertRange led fn) szNew . IntMap.insertWith RangeSet.union szNew RangeSet.empty)+          modify' $ case fitSub of+            Nothing -> over (eDB . unevaluated) (IntSet.delete sub . IntSet.delete subO)+            Just fs -> over (eDB . fitRangeDB) (removeRange sub fs . removeRange subO fs)+                     . over (eDB . sizeFitDB) (IntMap.adjust (removeRange subO fs . removeRange sub fs) szSub)       where         fitNew = (_fitness . _info) newC         fitLed = (_fitness . _info) ledC@@ -211,113 +305,93 @@         szSub  = (_size . _info) subC  -- | modify an e-class, e.g., add constant e-node and prune non-leaves-modifyEClass :: Monad m => CostFun -> EClassId -> EGraphST m EClassId+modifyEClass :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId modifyEClass costFun ecId =   do ec <- getEClass ecId-     -- let term = filter isTerm (Set.toList $ _eNodes ec)      case (_consts . _info) ec of-       ConstVal x -> do-         let en = Const x-         c <- calculateCost costFun en-         let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en }-         maybeEid <- gets ((Map.!? en) . _eNodeToEClass)-         modify' $ over eClass (IntMap.insert ecId ec{_eNodes = Set.singleton (encodeEnode en) , _info = infoEc})-         when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId)-         case maybeEid of-           Nothing   -> pure ecId-           Just eid' -> merge costFun eid' ecId+       ConstVal x ->+        do let en = EConst x+           c <- calculateCost costFun en+           let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en }+           maybeEid <- lookupNode en+           -- write through 'insertClass' (a paged store keeps the authoritative page)+           insertClass ec{ _eNodes = Set.singleton en, _info = infoEc }+           when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.insert ecId)+           case maybeEid of+             Nothing   -> pure ecId+             Just eid' -> merge costFun eid' ecId -       ParamIx x -> do-         let en = Param x-         c <- calculateCost costFun en-         ens <- gets (_eNodes . (IntMap.! ecId) . _eClass)-         let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en }-         maybeEid <- gets ((Map.!? en) . _eNodeToEClass)-         modify' $ over eClass (IntMap.insert ecId ec{_eNodes = Set.insert (encodeEnode en) (_eNodes ec), _info = infoEc})-         when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId)-         -- TODO: what happen to the orphans?-         case maybeEid of-           Nothing   -> pure ecId-           Just eid' -> merge costFun eid' ecId+       ParamIx x ->+         do let en = EParam x+            c <- calculateCost costFun en+            let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en }+            maybeEid <- lookupNode en+            insertClass ec{ _eNodes = Set.insert en (_eNodes ec), _info = infoEc }+            when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.insert ecId)+            case maybeEid of+              Nothing   -> pure ecId+              Just eid' -> merge costFun eid' ecId         _ -> pure ecId    where-    isTerm (Var _)   = True-    isTerm (Const _) = True-    isTerm (Param _) = True-    isTerm _         = False+    isTerm (EVar _)   = True+    isTerm (EConst _) = True+    isTerm (EParam _) = True+    isTerm _          = False -    toConst (Param ix) = ParamIx ix-    toConst (Const x)  = ConstVal x-    toConst _          = NotConst+    toConst (EParam ix) = ParamIx ix+    toConst (EConst x)  = ConstVal x+    toConst _           = NotConst  -- * DB --- | `createDB` creates a database of patterns from an e-graph--- it simply calls addToDB for every pair (e-node, e-class id) from--- the e-graph.-createDB :: Monad m => EGraphST m DB-createDB = do modify' $ over (eDB . patDB) (const Map.empty)-              ecls <- gets (Map.toList . _eNodeToEClass)-              mapM_ (uncurry addToDB) ecls-              gets (_patDB . _eDB)-{-# INLINE createDB #-}--createDBBest :: Monad m => EGraphST m DB-createDBBest = do modify' $ over (eDB . patDB) (const Map.empty)-                  ecls <- gets (Prelude.map (\(eId, ec) -> (_best (_info ec), eId)) . IntMap.toList . _eClass)-                  mapM_ (uncurry addToDB) ecls-                  gets (_patDB . _eDB)- -- | `addToDB` adds an e-node and e-class id to the database-addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB ()+addToDB :: (ClassStore m, HasCallStack) => ENode -> EClassId -> EGraphST m () -- State DB () addToDB enode' eid = do   eid' <- canonical eid-  isConst <- gets (_consts . _info . (IntMap.! eid') . _eClass)+  ec <- getEClass eid'+  let isConst = _consts . _info $ ec   let enode = case isConst of-                ConstVal x -> Const x-                ParamIx  x -> Param x+                ConstVal x -> EConst x+                ParamIx  x -> EParam x                 _          -> enode'-  let ids = eid : childrenOf enode -- we will add the e-class id and the children ids-      op  = getOperator enode    -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB-  trie <- gets ((Map.!? op) . _patDB . _eDB)       -- gets the entry for op, if it exists+  let ids = eid : eChildren enode -- we will add the e-class id and the children ids+      op  = eOpKey enode    -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB+  trie <- gets (Map.lookup op . _patDB . _eDB)   case populate trie ids of      -- populates the trie     Nothing -> pure ()     Just t  -> modify' $ over (eDB . patDB) (Map.insert op t) -- if something was created, insert back into the DB+  recordNode enode eid            -- register the node for the streaming matcher's source {-# INLINE addToDB #-}  -- | Populates an IntTrie with a sequence of e-class ids populate :: Maybe IntTrie -> [EClassId] -> Maybe IntTrie populate _ []         = Nothing--- if it is a new entry, simply add the ids sequentially populate Nothing eids = foldr f Nothing eids   where     f :: EClassId -> Maybe IntTrie -> Maybe IntTrie-    f eid (Just t) = Just $ trie eid (IntMap.singleton eid t)-    f eid Nothing  = Just $ trie eid IntMap.empty--- if the entry already exists, insert the new key--- and populate the next child entry recursivelly-populate (Just tId) (eid:eids) = let keys     = Set.insert eid (_keys tId)-                                     nextTrie = _trie tId IntMap.!? eid-                                     val      = fromMaybe (trie eid IntMap.empty) $ populate nextTrie eids-                                  in Just $ IntTrie keys (IntMap.insert eid val (_trie tId))+    f eid (Just t) = Just $ IntTrie (IntMap.singleton eid t)+    f eid Nothing  = Just $ IntTrie (IntMap.singleton eid (IntTrie IntMap.empty))+populate (Just tId) (eid:eids) = let nextTrie = IntMap.lookup eid (_trie tId)+                                     val      = fromMaybe (IntTrie IntMap.empty) $ populate nextTrie eids+                                  in Just $ IntTrie (IntMap.insert eid val (_trie tId)) {-# INLINE populate #-} -canonizeMap :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m (Map ClassOrVar ClassOrVar, ClassOrVar)-canonizeMap (subst, cv) = (,cv) <$> traverse g subst -- Map.fromList <$> traverse f (Map.toList subst)+canonizeMap :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m (Subst, ClassOrVar)+canonizeMap (subst, cv) = (,cv) <$> traverse g subst   where-    g :: Monad m => ClassOrVar -> EGraphST m ClassOrVar-    g (Left e2) = Left <$> canonical e2-    g e2        = pure e2--    f :: Monad m => (ClassOrVar, ClassOrVar) -> EGraphST m (ClassOrVar, ClassOrVar)-    f (e1, Left e2) = do e2' <- canonical e2-                         pure (e1, Left e2')-    f (e1, e2)      = pure (e1, e2)+    g :: ClassStore m => SubVal -> EGraphST m SubVal+    g (SVOne e2)  = SVOne <$> canonOne e2+    g (SVMap m)   = SVMap . IntMap.fromListWith (+) <$> mapM (\(e2, n) -> do+                       e2' <- canonOne (Left e2)+                       pure (getInt e2', n)) (IntMap.toList m)+    canonOne :: ClassStore m => ClassOrVar -> EGraphST m ClassOrVar+    canonOne (Left e2) = Left <$> canonical e2+    canonOne e2        = pure e2 {-# INLINE canonizeMap #-} -applyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m ()+applyMatch :: (ClassStore m, HasCallStack) => CostFun -> Rule -> (Subst, ClassOrVar) -> EGraphST m () applyMatch costFun rule match' =   do let conds = getConditions rule      match       <- canonizeMap match'@@ -329,28 +403,16 @@           pure () {-# INLINE applyMatch #-} -applyMergeOnlyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m ()-applyMergeOnlyMatch costFun rule match' =-  do let conds = getConditions rule-     match       <- canonizeMap match'-     validHeight <- isValidHeight match-     validConds  <- mapM (`isValidConditions` match) conds-     when (validHeight && and validConds) $-       do maybe_eid <- classOfENode costFun (fst match) (target rule)-          case maybe_eid of-            Nothing  -> pure ()-            Just eid -> do merge costFun (getInt (snd match)) eid-                           pure ()-{-# INLINE applyMergeOnlyMatch #-}- -- | gets the e-node of the target of the rule -- TODO: add consts and modify-classOfENode :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m (Maybe EClassId)-classOfENode costFun subst (VarPat c)     = do let maybeEid = getInt <$> subst Map.!? Right (fromEnum c)+classOfENode :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m (Maybe EClassId)+classOfENode costFun subst (VarPat c)     = do let maybeEid = case Map.lookup (Right (fromEnum c)) subst of+                                                                  Just (SVOne v) -> Just v+                                                                  _              -> Nothing                                                case maybeEid of                                                  Nothing  -> pure Nothing-                                                 Just eid -> Just <$> canonical eid-classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (Const x)+                                                 Just eid -> Just <$> canonical (getInt eid)+classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (EConst x) classOfENode costFun subst (Fixed target) = do newChildren <- mapM (classOfENode costFun subst) (getElems target)                                                case sequence newChildren of                                                  Nothing -> pure Nothing@@ -358,42 +420,111 @@                                                                cs' <- mapM canonical cs                                                                areConsts <- mapM isConst cs'                                                                if and areConsts-                                                                 then do eid <- add costFun new_enode+                                                                 then do eid <- addTree costFun new_enode                                                                          rebuild costFun -- eid new_enode                                                                          pure (Just eid)-                                                                 else gets ((Map.!? new_enode) . _eNodeToEClass)+                                                                 else do en <- toENode new_enode+                                                                         en' <- canonize en+                                                                         gets (HashMap.lookup en' . _eNodeToEClass)+classOfENode _ _ (NAry _ _) = error "classOfENode: n-ary pattern unsupported"+classOfENode _ _ Hole       = error "classOfENode: Hole is only valid in MapP targets" {-# INLINE classOfENode #-}  -- | adds the target of the rule into the e-graph-reprPrat :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m EClassId-reprPrat costFun subst (VarPat c)     = canonical $ getInt $ subst Map.! Right (fromEnum c)+reprPrat :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m EClassId+reprPrat costFun subst (VarPat c)     = do+    let k = Right (fromEnum c)+    v <- case Map.lookup k subst of+           Nothing -> error $ "REPRPRAT_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst)+           Just (SVOne x) -> pure x+           Just (SVMap _) -> error $ "REPRPRAT_REST_AS_SINGLE var=" <> show (fromEnum c)+    canonical $ getInt v reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target)-                                           add costFun (replaceChildren newChildren target)+                                           addTree costFun (replaceChildren newChildren target)+reprPrat costFun subst Hole = error "REPRPRAT_HOLE: Hole must be filled by MapP"+reprPrat costFun subst (NAry op ncs) = do+    m <- IntMap.unionsWith (+) <$> mapM (childEidM costFun subst) ncs+    case IntMap.toList m of+      []        -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1)))+      [(c, 1)]  -> canonical c+      _         -> do en <- mkENaryM op m+                      add costFun en {-# INLINE reprPrat #-} -isValidHeight :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool+-- | Adds a single child of an n-ary target pattern to the e-graph, returning+-- its contribution as a canonical multiset (so 'Rest' children carry their+-- 'IntMap' straight through without expansion).+childEidM :: (ClassStore m, HasCallStack) => CostFun -> Subst -> NChild -> EGraphST m (IntMap Int)+childEidM costFun subst (Ch p)     = (`IntMap.singleton` 1) <$> reprPrat costFun subst p+childEidM costFun subst (Rest c)   = restEidsM subst c+childEidM costFun subst (MapP p c) = do+  es <- restEids subst c+  ms <- forM es $ \e -> reprMapP costFun subst e p+  pure (imFromList ms)+{-# INLINE childEidM #-}++-- | The e-class ids bound to a rest variable, as a canonical multiset.+restEidsM :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m (IntMap Int)+restEidsM subst c = do+  let k = Right (fromEnum c)+  case Map.lookup k subst of+    Just (SVMap m) -> pure m+    Just (SVOne _) -> error $ "REPRPRAT_SINGLE_AS_REST var=" <> show (fromEnum c)+    Nothing        -> error $ "REPRPRAT_MISSING_REST var=" <> show (fromEnum c)+{-# INLINE restEidsM #-}++-- | The e-class ids bound to a rest variable, expanded one entry per+-- occurrence (used by 'MapP', which needs to instantiate per child).+restEids :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m [EClassId]+restEids subst c = expandedList <$> restEidsM subst c+{-# INLINE restEids #-}++-- | Build the target of a pattern where every `Hole` is filled with the+-- e-class `e` (used by 'MapP').+reprMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> Pattern -> EGraphST m EClassId+reprMapP costFun subst e Hole = canonical e+reprMapP costFun subst e (VarPat c) = reprPrat costFun subst (VarPat c)+reprMapP costFun subst e (Fixed target) = do+  newChildren <- mapM (reprMapP costFun subst e) (getElems target)+  addTree costFun (replaceChildren newChildren target)+reprMapP costFun subst e (NAry op ncs) = do+  m <- IntMap.unionsWith (+) <$> mapM (childMapP costFun subst e) ncs+  case IntMap.toList m of+    []   -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1)))+    [(c, 1)] -> canonical c+    _    -> do en <- mkENaryM op m+               add costFun en+{-# INLINE reprMapP #-}++-- | A single child of an n-ary pattern inside a 'MapP' function.+childMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> NChild -> EGraphST m (IntMap Int)+childMapP costFun subst e (Ch p)     = (`IntMap.singleton` 1) <$> reprMapP costFun subst e p+childMapP costFun subst e (Rest c)   = restEidsM subst c+childMapP costFun subst e (MapP _ _) = error "nested MapP unsupported"+{-# INLINE childMapP #-}++isValidHeight :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m Bool isValidHeight match = do-    h <- case snd match of-           Left ec -> do ec' <- canonical ec-                         gets (_height . (IntMap.! ec') . _eClass)-           Right _ -> pure 0-    pure $ h < 15+      h <- case snd match of+             Left ec -> _height <$> getEClass ec+             Right _ -> pure 0+      pure $ h < 15 {-# INLINE isValidHeight #-}  -- | returns `True` if the condition of a rule is valid for that match-isValidConditions :: Monad m => Condition -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool-isValidConditions cond match = gets $ cond (fst match)+isValidConditions :: ClassStore m => Condition -> (Subst, ClassOrVar) -> EGraphST m Bool+isValidConditions (Condition f) match = f (fst match) {-# INLINE isValidConditions #-}  -- * Tree to e-graph conversion and utility functions  -- | Creates an e-graph from an expression tree-fromTree :: Monad m => CostFun -> Fix SRTree -> EGraphST m EClassId-fromTree costFun = cataM sequence (add costFun)+fromTree :: (ClassStore m, HasCallStack) => CostFun -> Fix SRTree -> EGraphST m EClassId+fromTree costFun = cataM sequence (addTree costFun) {-# INLINE fromTree #-}  -- | Builds an e-graph from multiple independent trees-fromTrees :: Monad m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId]+fromTrees :: ClassStore m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId] fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) [] {-# INLINE fromTrees #-} @@ -403,119 +534,93 @@ countParamsUniqEg eg rt = countParamsUniq . runIdentity $ getBestExpr rt `evalStateT` eg  --- | gets the best expression given the default cost function-getBestExpr :: Monad m => EClassId -> EGraphST m (Fix SRTree)-getBestExpr eid = do eid' <- canonical eid-                     best <- gets (_best . _info . (IntMap.! eid') . _eClass)-                     childs <- mapM getBestExpr $ childrenOf best-                     pure . Fix $ replaceChildren childs best-{-# INLINE getBestExpr #-}--getBestENode eid = do eid' <- canonical eid-                      gets (_best . _info . (IntMap.! eid') . _eClass)+getBestENode eid = (_best . _info) <$> getEClass eid {-# INLINE getBestENode #-}  -- | returns one expression rooted at e-class `eId` -- TODO: avoid loopings-getExpressionFrom :: Monad m => EClassId -> EGraphST m (Fix SRTree)+getExpressionFrom :: ClassStore m => EClassId -> EGraphST m (Fix SRTree) getExpressionFrom eId' = do-    eId <- canonical eId'-    nodes <- gets (Set.map decodeEnode . _eNodes . (IntMap.! eId) . _eClass)-    let hasTerm = any isTerm nodes-        cands   = if hasTerm then filter isTerm (Set.toList nodes) else Set.toList nodes--    Fix <$> case head $ Set.toList nodes of-      Bin op l r -> Bin op <$> getExpressionFrom l <*> getExpressionFrom r-      Uni f t    -> Uni f <$> getExpressionFrom t-      Var ix     -> pure $ Var ix-      Const x    -> pure $ Const x-      Param ix   -> pure $ Param ix-  where-    isTerm (Var _) = True-    isTerm (Const _) = True-    isTerm (Param _) = True-    isTerm _ = False+    nodes <- _eNodes <$> getEClass eId'+    case Set.toList nodes of+      (n:_) -> case n of+        EVar ix     -> pure $ Fix $ Var ix+        EParam ix   -> pure $ Fix $ Param ix+        EConst x    -> pure $ Fix $ Const x+        EUni f t    -> Fix . Uni f <$> getExpressionFrom t+        EBin op l r -> Fix <$> (Bin op <$> getExpressionFrom l <*> getExpressionFrom r)+        ENAry op xs -> naryTree op <$> mapM getExpressionFrom (expandedList xs)+      [] -> error "getExpressionFrom: empty eclass" {-# INLINE getExpressionFrom #-}  -- | returns all expressions rooted at e-class `eId` -- TODO: check for infinite list-getAllExpressionsFrom :: Monad m => EClassId -> EGraphST m [Fix SRTree]+getAllExpressionsFrom :: ClassStore m => EClassId -> EGraphST m [Fix SRTree] getAllExpressionsFrom eId' = do-  eId <- canonical eId'-  nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)-  let cands  = filter isTerm nodes-  concat <$> go nodes-  --if null cands-  --   then concat <$> go nodes-  --   else pure [toTree $ head cands]+  nodes <- Set.toList . _eNodes <$> getEClass eId'+  go nodes   where-    isTerm (Var _) = True-    isTerm (Const _) = True-    isTerm (Param _) = True-    isTerm _ = False-    toTree (Var ix) = Fix $ Var ix-    toTree (Const x) = Fix $ Const x-    toTree (Param ix) = Fix $ Param ix-    toTree _ = undefined-     go []     = pure []     go (n:ns) = do-        t <- Prelude.map Fix <$> case n of-                Bin op l r -> do l' <- getAllExpressionsFrom l-                                 r' <- getAllExpressionsFrom r-                                 pure $ [Bin op li ri | li <- l', ri <- r']-                Uni f t    -> Prelude.map (Uni f) <$> getAllExpressionsFrom t-                Var ix     -> pure [Var ix]-                Const x    -> pure [Const x]-                Param ix   -> pure [Param ix]+        t <- case n of+                EVar ix     -> pure [Fix $ Var ix]+                EParam ix   -> pure [Fix $ Param ix]+                EConst x    -> pure [Fix $ Const x]+                EUni f t    -> Prelude.map (Fix . Uni f) <$> getAllExpressionsFrom t+                EBin op l r -> do l' <- getAllExpressionsFrom l+                                  r' <- getAllExpressionsFrom r+                                  pure $ [Fix $ Bin op li ri | li <- l', ri <- r']+                ENAry op xs -> do ts <- mapM getAllExpressionsFrom (expandedList xs)+                                  pure [ naryTree op comb | comb <- sequence ts ]         ts <- go ns-        pure (t:ts)+        pure (t ++ ts) {-# INLINE getAllExpressionsFrom #-} -getNExpressionsFrom :: Monad m => Int -> EClassId -> EGraphST m [Fix SRTree]+getNExpressionsFrom :: ClassStore m => Int -> EClassId -> EGraphST m [Fix SRTree] getNExpressionsFrom n eId' = getNExpressionsFrom' n 15 eId'  -getNExpressionsFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree]+getNExpressionsFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree] getNExpressionsFrom' _ 0 _ = pure [] getNExpressionsFrom' n d eId' = do-  eId <- canonical eId'-  nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)+  nodes <- Set.toList . _eNodes <$> getEClass eId'   (concat <$> go n d nodes)   where-    isTerm (Var _) = True-    isTerm (Const _) = True-    isTerm (Param _) = True+    isTerm (EVar _) = True+    isTerm (EConst _) = True+    isTerm (EParam _) = True     isTerm _ = False-    toTree (Var ix) = Fix $ Var ix-    toTree (Const x) = Fix $ Const x-    toTree (Param ix) = Fix $ Param ix+    toTree (EVar ix) = Fix $ Var ix+    toTree (EConst x) = Fix $ Const x+    toTree (EParam ix) = Fix $ Param ix     toTree _ = undefined      go n' _ []     = pure []     go n' 0 ts     = pure []     go n' d (node:ns) = do-        tt <- Prelude.map Fix <$> case node of-                Bin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l-                                 r' <- getNExpressionsFrom' n' (d-1) r-                                 pure $ Prelude.take n [Bin op li ri | li <- l', ri <- r']-                Uni f t    -> Prelude.map (Uni f) <$> getNExpressionsFrom' n' (d-1) t-                Var ix     -> pure [Var ix]-                Const x    -> pure [Const x]-                Param ix   -> pure [Param ix]+        tt <- case node of+                EVar ix     -> pure [Fix $ Var ix]+                EParam ix   -> pure [Fix $ Param ix]+                EConst x    -> pure [Fix $ Const x]+                EUni f t    -> Prelude.map (Fix . Uni f) <$> getNExpressionsFrom' n' (d-1) t+                EBin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l+                                  r' <- getNExpressionsFrom' n' (d-1) r+                                  pure $ Prelude.take n [Fix $ Bin op li ri | li <- l', ri <- r']+                ENAry op xs -> do ts <- mapM (getNExpressionsFrom' n' (d-1)) (expandedList xs)+                                  pure $ Prelude.take n [ naryTree op comb | comb <- sequence ts ]         let n'' = n' - length tt         if n'' <= 0           then pure [tt]           else do ts <- go n'' (d-1) ns                   pure (tt:ts) -getNEclassFrom :: Monad m => Int -> EClassId -> EGraphST m [[EClassId]]+getNEclassFrom :: ClassStore m => Int -> EClassId -> EGraphST m [[EClassId]] getNEclassFrom n eid = getNEclassFrom' n 15 eid -getNEclassFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [[EClassId]]+getNEclassFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [[EClassId]] getNEclassFrom' _ 0 _ = pure [] getNEclassFrom' n d eId' = do   eId <- canonical eId'-  nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)+  nodes <- Set.toList . _eNodes <$> getEClass eId'   (Prelude.map (eId:) <$> go n d nodes)   where     --go :: Int -> Int -> [ENode] -> EGraphST m [[EClassId]]@@ -523,13 +628,15 @@     go n' 0 ts     = pure []     go n' d (node:ns) = do         tt <- case node of-                Bin op l r -> do l' <- getNEclassFrom' n' (d-1) l-                                 r' <- getNEclassFrom' n' (d-1) r-                                 pure $ Prelude.take n [li <> ri | li <- l', ri <- r']-                Uni f t    -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]]-                Var ix     -> pure [[]]-                Const x    -> pure [[]]-                Param ix   -> pure [[]]+                EBin op l r -> do l' <- getNEclassFrom' n' (d-1) l+                                  r' <- getNEclassFrom' n' (d-1) r+                                  pure $ Prelude.take n [li <> ri | li <- l', ri <- r']+                ENAry op xs -> do ts <- mapM (getNEclassFrom' n' (d-1)) xs+                                  pure $ Prelude.take n [ concat comb | comb <- sequence ts ]+                EUni f t    -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]]+                EVar ix     -> pure [[]]+                EConst x    -> pure [[]]+                EParam ix   -> pure [[]]         pure tt         --let n'' = n' - length tt         --if n'' <= 0@@ -537,21 +644,21 @@         --  else do ts <- go n'' (d-1) ns         --          pure (tt:ts) -getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildEClasses :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildEClasses eId' = do   eId <- canonical eId'   IntSet.toList <$> go [eId] IntSet.empty    where     hasNoTerminal :: [ENode] -> Bool-    hasNoTerminal = all (not . null . childrenOf) -    getNodes :: Monad m => EClassId -> EGraphST m [ENode]-    getNodes n = gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)+    hasNoTerminal = all (not . null . eChildren) +    getNodes :: ClassStore m => EClassId -> EGraphST m [ENode]+    getNodes n = Set.toList . _eNodes <$> getEClass n -    go :: Monad m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet+    go :: ClassStore m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet     go [] visited = pure visited     go queue visited = do -        nodes <- concatMap childrenOf . concat . filter hasNoTerminal <$> mapM getNodes queue+        nodes <- concatMap eChildren . concat . filter hasNoTerminal <$> mapM getNodes queue         eids <- filter (\e -> e `IntSet.notMember` visited) <$> (mapM canonical nodes)         go eids (visited `IntSet.union` IntSet.fromList queue)             {-@@ -565,31 +672,27 @@                         -} {-# INLINE getAllChildEClasses #-} -getAllChildBestEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildBestEClasses :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildBestEClasses eId' = do-  eId <- canonical eId'-  nub <$> go eId-+  IntSet.toList <$> go IntSet.empty eId'   where-    go :: Monad m => Int -> EGraphST m [Int]-    go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass)-              let hasTerminal = (null . childrenOf) node-              eids <- mapM canonical $ childrenOf node-              if hasTerminal-                then pure [n]-                else do eids' <- mapM go eids-                        pure ((n : eids) <> concat eids')+    go :: ClassStore m => IntSet.IntSet -> EClassId -> EGraphST m IntSet.IntSet+    go acc n+      | IntSet.member n acc = pure acc+      | otherwise = do+          let acc' = IntSet.insert n acc+          node <- (_best . _info) <$> getEClass n+          eids <- mapM canonical $ eChildren node+          foldM go acc' eids -getAllChildBestEClassesRep :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildBestEClassesRep :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildBestEClassesRep eId' = do-  eId <- canonical eId'-  go eId-+  go eId'   where-    go :: Monad m => Int -> EGraphST m [Int]-    go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass)-              let hasTerminal = (null . childrenOf) node-              eids <- mapM canonical $ childrenOf node+    go :: ClassStore m => EClassId -> EGraphST m [EClassId]+    go n = do node <- (_best . _info) <$> getEClass n+              let hasTerminal = (null . eChildren) node+              eids <- mapM canonical $ eChildren node               if hasTerminal                 then pure [n]                 else do eids' <- mapM go eids@@ -598,43 +701,43 @@ -- | returns a random expression rooted at e-class `eId` getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree) getRndExpressionFrom eId' = do-    eId <- canonical eId'-    nodes <- gets (Set.toList . _eNodes . (IntMap.! eId) . _eClass)+    nodes <- Set.toList . _eNodes <$> getEClass eId'     n <- lift $ randomFrom nodes-    Fix <$> case decodeEnode n of-              Bin op l r -> Bin op <$> getRndExpressionFrom l <*> getRndExpressionFrom r-              Uni f t    -> Uni f <$> getRndExpressionFrom t-              Var ix     -> pure $ Var ix-              Const x    -> pure $ Const x-              Param ix   -> pure $ Param ix+    case n of+      EUni f t    -> Fix . Uni f <$> getRndExpressionFrom t+      EBin op l r -> Fix <$> (Bin op <$> getRndExpressionFrom l <*> getRndExpressionFrom r)+      ENAry op xs -> naryTree op <$> mapM getRndExpressionFrom (expandedList xs)+      EVar ix     -> pure $ Fix $ Var ix+      EConst x    -> pure $ Fix $ Const x+      EParam ix   -> pure $ Fix $ Param ix   where     randomRange rng = state (randomR rng)     randomFrom xs   = do n <- randomRange (0, length xs - 1)                          pure $ xs !! n {-# INLINE getRndExpressionFrom #-} -cleanMaps :: Monad m => EGraphST m ()+cleanMaps :: ClassStore m => EGraphST m () cleanMaps = do-  enode2eclass <- gets _eNodeToEClass-  entries <- forM (Map.toList enode2eclass) $ \(k,v) -> do-    k' <- canonize k-    v' <- canonical v-    pure (k',v')-  let enode2eclass' = Map.fromList entries-  eclassMap <- gets _eClass-  entries' <- forM (IntMap.toList eclassMap) $ \(k,v) -> do-    k' <- canonical k-    pure $ if k==k' then (Just (k,v)) else Nothing-  let eclassMap' = IntMap.fromList (catMaybes entries')-  canon <- gets _canonicalMap-  entries'' <- forM (IntMap.toList canon) $ \(k,v) -> do-    pure $ if k==v then Just (k,v) else Nothing-  let canon' = IntMap.fromList (catMaybes entries'')-  eDB' <- gets _eDB-  put $ EGraph canon enode2eclass' eclassMap' eDB'-  forceState+  hasStore <- gets (isJust . _classStore)+  if hasStore+    -- the paged store is authoritative for both node->class and canonical+    -- lookups, so the bounded resident caches are simply reset (an O(n) rebuild+    -- of an unbounded map would defeat the out-of-core goal).+    then modify' $ \eg -> eg { _eNodeToEClass = HashMap.empty+                             , _canonicalMap = IntMap.empty+                             , _eClass = IntMap.empty }+    else do+      enode2eclass <- gets _eNodeToEClass+      entries <- forM (HashMap.toList enode2eclass) $ \(k,v) -> do+        k' <- canonize k+        v' <- canonical v+        pure (k',v')+      let enode2eclass' = HashMap.fromList entries+      eclassMap <- gets _eClass+      entries' <- forM (IntMap.toList eclassMap) $ \(k,v) -> do+        k' <- canonical k+        pure $ if k==k' then (Just (k,v)) else Nothing+      let eclassMap' = IntMap.fromList (catMaybes entries')+      modify' $ \eg -> eg { _eNodeToEClass = enode2eclass'+                          , _eClass = eclassMap' } {-# INLINE cleanMaps #-}--forceState :: Monad m => StateT s m ()-forceState = get >>= \ !_ -> return ()-{-# INLINE forceState #-}

@@ -1,5 +1,6 @@ {-# LANGUAGE TupleSections #-} {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE RankNTypes #-} ----------------------------------------------------------------------------- -- | -- Module      :  Algorithm.EqSat.EqSatDB@@ -19,26 +20,36 @@ import Control.Lens ( over ) import Control.Monad (when, foldM, forM) import Control.Monad.State-import Data.IntMap (IntMap)-import qualified Data.IntMap as IntMap-import Data.List (intercalate, nub, sortBy)+import GHC.Stack (HasCallStack)+import Data.IntMap.Strict (IntMap)+import qualified Data.IntMap.Strict as IntMap import Data.Map (Map) import qualified Data.Map as Map+import Data.List (sortBy) import Data.Maybe (fromMaybe) import Data.Ord (comparing) import Data.SRTree---import Data.Set (Set) import Data.HashSet (HashSet) import qualified Data.HashSet as Set+import qualified Data.Set as RangeSet import Data.String (IsString (..)) import Data.SRTree.Recursion (cata)+import Text.Read (readMaybe) -import Debug.Trace --- A Pattern is either a fixed-point of a tree or an--- index to a pattern variable. The pattern variable matches anything. -data Pattern = Fixed (SRTree Pattern) | VarPat Char deriving (Show, Eq, Ord) -- Fixed structure of a pattern or a variable that matches anything+-- A Pattern is either a fixed-point of a tree, an index to a pattern variable+-- (which matches anything), a hole (only used inside a 'MapP' target function),+-- or an n-ary Add/Mul pattern whose children are matched as a multiset.+data Pattern = Fixed (SRTree Pattern) | VarPat Char | Hole | NAry NOp [NChild]+  deriving (Show, Eq, Ord) +-- | A child of an n-ary pattern: a single child pattern ('Ch'), a rest+-- variable binding every remaining child of the node ('Rest'), or a+-- target-side map that splices one instantiation of a pattern (with its 'Hole'+-- filled) per child bound to a rest variable ('MapP').+data NChild = Ch Pattern | Rest Char | MapP Pattern Char+  deriving (Show, Eq, Ord)+ -- The instance for `IsString` for a `Pattern` is  -- valid only for a single letter char from a-zA-Z.  -- The patterns can be written as "x" + "y", for example,@@ -54,6 +65,8 @@     alg (Param ix) = if ix >= 100 then VarPat (toEnum $ ix - 100 + 65) else Fixed $ Param ix     alg (Var ix) = Fixed $ Var ix     alg (Const x) = Fixed $ Const x+    alg (Bin Add l r) = NAry EAdd [Ch l, Ch r]+    alg (Bin Mul l r) = NAry EMul [Ch l, Ch r]     alg (Bin op l r) = Fixed $ Bin op l r     alg (Uni f t) = Fixed $ Uni f t -- A rule is either a directional rule where pat1 can be replaced by pat2, a bidirectional rule @@ -73,32 +86,45 @@ -- A Query is a list of Atoms  type Query = [Atom] --- A `Condition` is a function that takes a substution map,--- an e-graph and returns whether the pattern attends the condition.-type Condition = Map ClassOrVar ClassOrVar -> EGraph -> Bool+-- | A `Condition` is a predicate over a match's substitution that runs inside+-- the e-graph monad so it can fetch e-class data through 'ClassStore' (which+-- streams from a paged store when the graph is out-of-core). The quantification+-- over the monad is intentional: the same condition works for any 'ClassStore'+-- instance, including the IO-backed paged store.+newtype Condition = Condition (forall m. ClassStore m => Subst -> EGraphST m Bool)  -- An Atom is composed of either an e-class id or pattern variable id -- and the tree that generated that pattern. Left is e-class id and Right is a VarPat. type ClassOrVar = Either EClassId Int data Atom = Atom ClassOrVar (SRTree ClassOrVar) deriving Show +-- | A substitution value: a single e-class (a matched pattern variable) or the+-- canonical multiset of e-class ids (a matched rest variable).+data SubVal = SVOne ClassOrVar | SVMap (IntMap Int) deriving Show++-- | Substitution map produced by matching a pattern.+type Subst = Map ClassOrVar SubVal+ unFixPat :: Pattern -> SRTree Pattern unFixPat (Fixed p) = p+unFixPat (VarPat _) = error "unFixPat: VarPat is not a fixed pattern"+unFixPat Hole       = error "unFixPat: Hole is not a fixed pattern"+unFixPat (NAry _ _) = error "unFixPat: NAry is not a fixed pattern" {-# INLINE unFixPat #-}   instance Num Pattern where-  l + r = Fixed $ Bin Add l r+  l + r = NAry EAdd [Ch l, Ch r]   {-# INLINE (+) #-}-  l - r = Fixed $ Bin Sub l r+  l - r = NAry EAdd [Ch l, Ch (negate r)]   {-# INLINE (-) #-}-  l * r = Fixed $ Bin Mul l r+  l * r = NAry EMul [Ch l, Ch r]   {-# INLINE (*) #-}    abs = Fixed . Uni Abs   {-# INLINE abs #-} -  negate t = Fixed (Const (-1)) * t+  negate t = NAry EMul [Ch (Fixed (Const (-1))), Ch t]   {-# INLINE negate #-}    signum t = case t of@@ -108,7 +134,7 @@   {-# INLINE fromInteger #-}  instance Fractional Pattern where-  l / r = Fixed $ Bin Div l r+  l / r = NAry EMul [Ch l, Ch (Fixed (Uni Recip r))]   {-# INLINE (/) #-}    fromRational = Fixed . Const . fromRational@@ -176,38 +202,347 @@ {-# INLINE cleanDB #-}  -- | Returns the substitution rules--- for every match of the pattern `source` inside the e-graph.-match :: Monad m => Pattern -> EGraphST m [(Map ClassOrVar ClassOrVar, ClassOrVar)]-match src = do-  let (q, root) = compileToQuery src     -- compile the source of the pattern into a query-  substs <- genericJoin q root               -- find the substituion rules for this pattern-  pure [(s, s Map.! root) | s <- substs, Map.size s > 0]+-- for every match of the pattern `source` inside the e-graph. This is the pure+-- matcher (no seen-set) used by user pattern queries; saturation uses+-- 'matchSaturated'.+match :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)]+match src = if hasNAry src+              then matchNAryWith Nothing src+              else do+                paged <- isPagedGraph+                if paged+                  then matchStreamCached Nothing src+                  else matchCachedWith Nothing (compileToQuery src) {-# INLINE match #-} --- | Returns a Query (list of atoms) of a pattern-compileToQuery :: Pattern -> (Query, ClassOrVar)-compileToQuery pat = evalState (processPat pat) 256 -- returns (atoms, root)+-- | Non-n-ary matching. The match's root e-class anchors it the same way the+-- n-ary matcher anchors one match per trie root, so it shares the same cheap+-- persistent mark-on-attempt seen-set ('_seenMatches', keyed by rule source ->+-- root class id): already-processed roots are skipped so the per-rule budget+-- advances to new matches across the scheduler's ban/unban cycles. Keying by+-- the root (an @O(1)@ class id) avoids serializing every substitution, which+-- would dominate on rules whose @genericJoin@ yields many matches. 'Nothing'+-- disables the seen-set (pure queries).+matchCachedWith :: ClassStore m => Maybe String -> (Query, [ClassOrVar], ClassOrVar) -> EGraphST m [(Subst, ClassOrVar)]+matchCachedWith mSk (q, vars, root) = do+  ss <- genericJoin q vars root+  seenSk <- case mSk of+    Nothing -> pure RangeSet.empty+    Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB)+  let rootOf s = case Map.lookup root s of+                   Just (SVOne (Left eid)) -> eid+                   _                       -> 0+      fresh = [ s | s <- ss+                  , Map.size s > 0+                  , maybe True (\_ -> not (RangeSet.member (show (rootOf s)) seenSk)) mSk ]+      taken = take ruleMatchBudget fresh+  case mSk of+    Just sk -> modify' $ over (eDB . seenMatches)+               (Map.insertWith RangeSet.union sk (RangeSet.fromList (map (show . rootOf) taken)))+    Nothing -> pure ()+  pure [ (s, case Map.lookup root s of+               Nothing -> error $ "MATCHCACHED_MISSING root=" <> show (getInt root) <> " substSize=" <> show (Map.size s)+               Just v  -> fromSVOne v)+       | s <- taken ]+{-# INLINE matchCachedWith #-}++-- | Saturation matching: consults/marks the persistent seen-set so each rule's+-- per-iteration budget advances to genuinely new matches across ban/unban.+matchSaturated :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)]+matchSaturated src = if hasNAry src+                       then matchNAryWith (Just (show src)) src+                       else do+                         paged <- isPagedGraph+                         if paged+                           then matchStreamCached (Just (show src)) src+                           else matchCachedWith (Just (show src)) (compileToQuery src)+{-# INLINE matchSaturated #-}++-- | True if the pattern (or a nested child) is an n-ary Add/Mul pattern.+hasNAry :: Pattern -> Bool+hasNAry (NAry _ _) = True+hasNAry (Fixed t)  = any hasNAry (getElems t)+hasNAry _          = False+{-# INLINE hasNAry #-}++-- | The operator trie key of the top-level pattern.+opOf :: Pattern -> SRTree ()+opOf (NAry EAdd _) = Bin Add () ()+opOf (NAry EMul _) = Bin Mul () ()+opOf (Fixed t)     = getOperator t+opOf _             = error "opOf: pattern has no operator"+{-# INLINE opOf #-}++-- | Matches an n-ary pattern against every root e-node of the operator trie.+-- A per-rule result budget ('ruleBudget') bounds the total number of matches+-- returned for one rule against one individual's nodes, and only the first+-- match per root e-class is kept, taming the O(k^2*m^2) backtracking of+-- Rest/Ch rules (e.g. factoring a common term out of a sum of products).+-- Keeping one match per root is sound: every returned match is genuine, and+-- the egraph merges the equivalent rewrites that further matches would apply,+-- so the rest of the root's matches are redundant work.+ruleBudget :: Int+ruleBudget = 64++-- | Cap on how many operator-trie root e-classes a single rule may visit per+-- match. 'ruleBudget' bounds the number of *results* returned, but a rule whose+-- matches are rare would otherwise still scan every root e-class in the trie+-- (every @+@/@*@ class in the graph), doing an expensive 'recursiveMatch' per+-- root -- which blows up on large graphs even though few matches result.+-- Capping root visits bounds the *search work* independently of the result+-- count. Sound: we only stop enumerating (fewer) genuine matches early.+ruleRootVisit :: Int+ruleRootVisit = 512++-- | Cap on how many matches a non-n-ary rule (the cached @genericJoin@ path)+-- may return per match. The n-ary matcher has 'ruleBudget'; give the cached+-- path a separate (larger) budget so a single rule cannot flood the iteration.+ruleMatchBudget :: Int+ruleMatchBudget = 1024++-- | Cap on how many operator-root e-classes the streaming cached matcher visits+-- per match, bounding the search work (and the page reads) independently of the+-- result count, exactly as 'ruleRootVisit' does for the n-ary matcher.+ruleMatchRootVisit :: Int+ruleMatchRootVisit = 2048++-- | Match an n-ary pattern against every root e-class of its operator trie.+--+-- A persistent per-source set of already-attempted roots ('_seenMatches') lets+-- the matcher skip roots it has already tried, so the per-rule result/search+-- budgets keep advancing to *new* roots across the scheduler's ban/unban cycles+-- instead of re-enumerating the same head of the trie (which starves the tail).+-- Roots are marked as attempted on the first try ('mark-on-attempt'), whether or+-- not they yielded a match, so a match that fails 'applyMatch' conditions is not+-- re-attempted every cycle.+matchNAryWith :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)]+matchNAryWith mSk src = do+  seen <- case mSk of+    Nothing -> pure RangeSet.empty+    Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB)+  -- skip already-attempted roots so the per-rule budget advances to new roots+  -- across the scheduler's ban/unban cycles (matches the trie path's semantics).+  let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ]+  roots <- streamRoots (opOf src) ruleRootVisit exclude+  go roots 0 0 []   where+    go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)]+    go [] _ _ acc = pure (reverse acc)+    go _ n _ acc | n >= ruleBudget = pure (reverse acc)+    go (_ : _) _ r acc | r >= ruleRootVisit = pure (reverse acc)+    go (eid : eids) n r acc = do+      -- mark-on-attempt: remember this root as tried for this rule source+      case mSk of+        Just sk -> modify' $ over (eDB . seenMatches)+                   (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid)))+        Nothing -> pure ()+      substs <- recursiveMatch src eid Map.empty+      let newMs = take 1 [ (s, Left eid) | s <- substs ]+      go eids (n + length newMs) (r + 1) (foldr (:) acc newMs)+{-# INLINE matchNAryWith #-}++-- | Streaming matcher for the cached (non-n-ary @genericJoin@) path on a paged+-- graph. Instead of enumerating candidates from the in-RAM @_patDB@ trie, it+-- streams the candidate root e-classes of the pattern's operator through+-- 'streamRoots' (bounded, skipping the already-attempted seen-set) and matches+-- each root incrementally with 'recursiveMatch' (which reads e-classes through+-- the paged store). This is the out-of-core analogue of 'matchCachedWith': the+-- resident/pure path keeps the optimized trie 'genericJoin', and only a paged+-- graph takes this route, so the matcher never builds an O(nodes) structure.+--+-- 'ruleMatchBudget' bounds the results and 'ruleMatchRootVisit' bounds the root+-- visits; the persistent mark-on-attempt seen-set makes each rule's budgets+-- advance to new roots across the scheduler's ban/unban cycles.+matchStreamCached :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)]+matchStreamCached mSk src = do+  seen <- case mSk of+    Nothing -> pure RangeSet.empty+    Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB)+  let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ]+  roots <- case opOfMay src of+             Just op  -> streamRoots op ruleMatchRootVisit exclude+             Nothing  -> pure []+  go roots 0 0 []+  where+    go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)]+    go [] _ _ acc = pure (reverse acc)+    go _ n _ acc | n >= ruleMatchBudget = pure (reverse acc)+    go (_ : _) _ r acc | r >= ruleMatchRootVisit = pure (reverse acc)+    go (eid : eids) n r acc = do+      case mSk of+        Just sk -> modify' $ over (eDB . seenMatches)+                   (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid)))+        Nothing -> pure ()+      substs <- recursiveMatch src eid Map.empty+      let newMs = take (ruleMatchBudget - n) [ (s, Left eid) | s <- substs ]+      go eids (n + length newMs) (r + 1) (foldr (:) acc newMs)+{-# INLINE matchStreamCached #-}++-- | The operator trie key of the top-level pattern, or @Nothing@ for a pattern+-- with no operator (e.g. a bare variable), which the streaming matcher treats+-- as matching nothing.+opOfMay :: Pattern -> Maybe (SRTree ())+opOfMay (NAry EAdd _) = Just (Bin Add () ())+opOfMay (NAry EMul _) = Just (Bin Mul () ())+opOfMay (Fixed t)     = Just (getOperator t)+opOfMay _             = Nothing+{-# INLINE opOfMay #-}++-- | Recursively match a pattern against the e-class `eid`, threading a+-- substitution map, returning every substitution that completes the match.+recursiveMatch :: ClassStore m => Pattern -> EClassId -> Subst -> EGraphST m [Subst]+recursiveMatch (VarPat c) eid subst =+  pure (bindVar subst (Right (fromEnum c)) eid)+recursiveMatch Hole _ subst = pure [subst]+recursiveMatch (Fixed t) eid subst = matchFixed t eid subst+recursiveMatch (NAry op ncs) eid subst = matchNAryNode op ncs eid subst+{-# INLINE recursiveMatch #-}++-- | Bind `v` to the e-class `eid`, enforcing that re-occurrences of `v` are+-- consistent.+bindVar :: Subst -> ClassOrVar -> EClassId -> [Subst]+bindVar subst v eid =+  case Map.lookup v subst of+    Just (SVOne e) | e == Left eid -> [subst]+    Just _                         -> []+    Nothing                        -> [Map.insert v (SVOne (Left eid)) subst]+{-# INLINE bindVar #-}++-- | Match a fixed tree pattern against the e-nodes of the e-class `eid`,+-- returning every substitution that completes the match across all candidate+-- e-nodes.+matchFixed :: ClassStore m => SRTree Pattern -> EClassId -> Subst -> EGraphST m [Subst]+matchFixed t eid subst = do+  ec <- getEClass eid+  let cands = [n | n <- Set.toList (_eNodes ec), eOpKey n == getOperator t]+  fmap concat $ forM cands $ \n -> matchChildren t subst n+  where+    matchChildren t s n = go (zip (getElems t) (enodeChildren n)) [s]+    go [] ss = pure ss+    go ((p, c) : ps) ss = do+      ms <- concat <$> mapM (\s -> recursiveMatch p c s) ss+      go ps ms+{-# INLINE matchFixed #-}++-- | The child e-class ids of an e-node, in canonical (sorted for ENAry) order.+enodeChildren :: ENode -> [EClassId]+enodeChildren (EUni _ t)   = [t]+enodeChildren (EBin _ l r) = [l, r]+enodeChildren (ENAry _ m)  = expandedList m+enodeChildren _            = []+{-# INLINE enodeChildren #-}++-- | Match an n-ary pattern node against the e-class `eid`: it must contain an+-- ENAry node of the given op, whose children are matched as a multiset. Every+-- ENAry node in the class is tried.+matchNAryNode :: ClassStore m => NOp -> [NChild] -> EClassId -> Subst -> EGraphST m [Subst]+matchNAryNode op ncs eid subst = do+  ec <- getEClass eid+  let nodes = [m | ENAry op' m <- Set.toList (_eNodes ec), op' == op]+  fmap concat $ forM nodes $ \m ->+    matchNChildren ncs m subst+{-# INLINE matchNAryNode #-}++-- | Match a sequence of n-ary children against a multiset of e-class ids.+-- Each 'Ch' consumes one matched child; a 'Rest' child consumes all remaining+-- children. Every multiset assignment is returned. Iterating over the distinct+-- child ids (the multiset's keys) is sound (duplicate copies only differ by+-- position, which 'decChild' already resolves) and avoids duplicate result+-- sets.+--+-- A per-call result budget ('matchCap') caps the number of substitutions+-- returned, bounding the O(k^2*m^2) backtracking of Rest/Ch rules such as+-- factoring a common term out of a sum of products. Sound: each result is a+-- genuine match; we merely stop enumerating once the budget is exhausted.+matchCap :: Int+matchCap = 64++matchNChildren :: ClassStore m => [NChild] -> IntMap Int -> Subst -> EGraphST m [Subst]+matchNChildren ncs children subst = reverse <$> goB ncs children subst matchCap+  where+    goB :: ClassStore m => [NChild] -> IntMap Int -> Subst -> Int -> EGraphST m [Subst]+    goB [] m s _+      | IntMap.null m = pure [s]+      | otherwise     = pure []+    goB (Rest c : ps) m s b = do+      let v = Right (fromEnum c)+      case Map.lookup v s of+        Just _  -> pure []  -- rest variable already bound+        Nothing -> goB ps IntMap.empty (Map.insert v (SVMap m) s) b+    goB (Ch p : ps) m s b+      | multiplicity m <= nCh ps = pure []  -- not enough children left+      | otherwise = goC (IntMap.keys m) 0 []+      where+        goC :: ClassStore m => [EClassId] -> Int -> [Subst] -> EGraphST m [Subst]+        goC [] _ acc = pure acc+        goC _ n acc | n >= b    = pure acc+        goC (c : cs) n acc = do+          ms <- recursiveMatch p c s+          goMs c ms cs n acc+        goMs :: ClassStore m => EClassId -> [Subst] -> [EClassId] -> Int -> [Subst] -> EGraphST m [Subst]+        goMs c [] cs n acc = goC cs n acc+        goMs c (s' : ms) cs n acc+          | n >= b     = pure acc+          | otherwise = do+              r <- goB ps (decChild c m) s' (b - n)+              let r' = take (b - n) r+                  n' = n + length r'+              goMs c ms cs n' (foldr (:) acc r')+    goB (MapP _ _ : _) _ _ _ = error "matchNChildren: MapP is only valid in targets"+{-# INLINE matchNChildren #-}++-- | Total number of children (counting multiplicities) in a multiset.+multiplicity :: IntMap Int -> Int+multiplicity = IntMap.foldr' (+) 0+{-# INLINE multiplicity #-}++-- | Remove one occurrence of `c` from the multiset (decrementing its+-- multiplicity, or dropping the key entirely when it reaches zero).+decChild :: Int -> IntMap Int -> IntMap Int+decChild c = IntMap.update (\n -> if n > 1 then Just (n - 1) else Nothing) c+{-# INLINE decChild #-}++-- | Number of 'Ch' patterns in a child pattern sequence (each consumes one+-- child, so at least this many children must remain).+nCh :: [NChild] -> Int+nCh = length . filter isCh+  where+    isCh (Ch _)   = True+    isCh _        = False+{-# INLINE nCh #-}++-- | Unwrap a single-e-class substitution value.+fromSVOne :: SubVal -> ClassOrVar+fromSVOne (SVOne v)    = v+fromSVOne (SVMap _)    = error "fromSVOne: expected a single e-class"+{-# INLINE fromSVOne #-}++-- | Returns a Query (list of atoms) of a pattern with pre-computed ordered vars+compileToQuery :: Pattern -> (Query, [ClassOrVar], ClassOrVar)+compileToQuery pat = (atoms, orderedVars atoms, root)+  where (atoms, root) = evalState (processPat pat) 256       -- creates the atoms of a pattern-      processPat :: Pattern -> State Int (Query, ClassOrVar)-      processPat (VarPat x)  = pure ([], Right $ fromEnum x)-      processPat (Fixed pat) = do-          -- get the next available var id and add as root-          v <- get-          let root = Right v-          -- updates the next available id-          modify (+1)-          -- recursivelly process the children of the pattern-          patChilds <- mapM processPat (getElems pat)-          -- create an atom composed of the-          -- root and the tree with the children-          -- replaced by the childs roots-          -- add the child atoms to the list-          let atoms = concatMap fst patChilds-              roots = map snd patChilds-              atom  = Atom root (replaceChildren roots pat)-              atoms' = atom:atoms-          pure (atoms', root)+        processPat :: Pattern -> State Int (Query, ClassOrVar)+        processPat (VarPat x)  = pure ([], Right $ fromEnum x)+        processPat (NAry _ _)  = error "compileToQuery: n-ary pattern (use matchNAry instead)"+        processPat Hole        = error "compileToQuery: Hole is only valid in MapP targets"+        processPat (Fixed pat) = do+            -- get the next available var id and add as root+            v <- get+            let root = Right v+            -- updates the next available id+            modify (+1)+            -- recursivelly process the children of the pattern+            patChilds <- mapM processPat (getElems pat)+            -- create an atom composed of the+            -- root and the tree with the children+            -- replaced by the childs roots+            -- add the child atoms to the list+            let atoms = concatMap fst patChilds+                roots = map snd patChilds+                atom  = Atom root (replaceChildren roots pat)+                atoms' = atom:atoms+            pure (atoms', root) {-# INLINE compileToQuery #-}  -- get the value from the Either Int Int@@ -226,63 +561,51 @@ -- | Creates the substituion map for -- the pattern variables for each one of the -- matched subgraph-genericJoin :: Monad m => Query -> ClassOrVar -> EGraphST m [Map ClassOrVar ClassOrVar]-genericJoin atoms root = do-  let vars = orderedVars atoms -- order the vars, starting with the most frequently occuring-  go atoms vars -- TODO: investigate why we need nub+genericJoin :: (ClassStore m, HasCallStack) => Query -> [ClassOrVar] -> ClassOrVar -> EGraphST m [Subst]+genericJoin atoms vars root = go atoms vars   where     -- for each variable     --   for each possible e-class id for that variable     --      replace the var id with this e-class id, and     --      recurse to find the possible matches for the next atom-    go :: Monad m => Query -> [ClassOrVar] -> EGraphST m [Map ClassOrVar ClassOrVar]+    go :: ClassStore m => Query -> [ClassOrVar] -> EGraphST m [Subst]     go atoms [] = pure [Map.empty] -- | _ <- atoms]     go atoms (x:vars) = do cIds1 <- domainX x atoms root                            maps <- forM cIds1 $ \classId -> do-                             map (Map.insert x classId) <$> go (updateVar x classId atoms) vars+                             map (Map.insert x (SVOne classId)) <$> go (updateVar x classId atoms) vars                            pure (concat maps) {-# INLINE genericJoin #-} -     -- [Map.insert x classId y | classId <- domainX db x atoms-     --                                           , y <- go (updateVar x classId atoms) vars]   -- | returns the e-class id for a certain variable that -- matches the pattern described by the atoms-domainX :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar]+domainX :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar] domainX var atoms root = do   let atoms' = filter (elemOfAtom var) atoms -- :: [ClassOrVar]  -- look only in the atoms with this var   map Left <$> intersectAtoms var atoms' root -- find the intersection of possible keys by each atom {-# INLINE domainX #-}-  --let ss = (map Left-  --                                $ intersectAtoms var db-  --                                $-  --                     in ss  -- | returns all e-class id that can matches this sequence of atoms-intersectAtoms :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId]+intersectAtoms :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId] intersectAtoms _ [] root = pure [] intersectAtoms var (a:atoms) root = do-  a0 <- go a-  Set.toList <$> (foldM (\acc atom -> Set.intersection acc <$> go atom) a0 atoms)+  a0 <- toCanon =<< go a+  Set.toList <$> (foldM (\acc atom -> do+    res <- go atom+    Set.intersection acc <$> toCanon res) a0 atoms)   where-      -- canonize everything except the root for consistency-      -- doing this here prevents traversing the map again       toCanon x = if var==root                      then pure x                      else Set.fromList <$> (mapM canonical $ Set.toList x) -      go (Atom r t) = do-        let op = getOperator t-        mTrie <- gets ((Map.!? op) . _patDB . _eDB)-        case mTrie of-          Just trie -> pure (fromMaybe Set.empty $ intersectTries var Map.empty trie (r:getElems t))-          Nothing   -> pure Set.empty-          -- TODO: remove FlexibleContexts-        --if op `Map.member` db -- if the e-graph contains the operator-                               -- try to find an intersection of the tries that matches each atom of the pattern-        --  then-        --  else pure Set.empty+      go (Atom r t) =+        do let op = getOperator t+           mTrie <- gets ((Map.!? op) . _patDB . _eDB)+           case mTrie of+             Just trie -> pure (fromMaybe Set.empty $ intersectTries var IntMap.empty trie (r:getElems t))+             Nothing   -> pure Set.empty+ {-# INLINE intersectAtoms #-}  -- | searches for the intersection of e-class ids that@@ -294,40 +617,26 @@ -- trie is the current trie of the pattern -- (i:ids) sequence of root : children of the atom to investigate -- NOTE: it must be Maybe Set to differentiate between empty set and no answer-intersectTries :: ClassOrVar -> Map ClassOrVar EClassId -> IntTrie -> [ClassOrVar] -> Maybe (HashSet EClassId)+intersectTries :: ClassOrVar -> IntMap EClassId -> IntTrie -> [ClassOrVar] -> Maybe (HashSet EClassId) intersectTries var xs trie [] = Just Set.empty intersectTries var xs trie (i:ids) =     case i of-      Left x  -> if x `Set.member` _keys trie-                    -- if the current investigated id is an e-class id and-                    -- it is one of the keys of the trie...-                    -- ..try to match the next id with the next trie-                    then intersectTries var xs (_trie trie IntMap.! x) ids-                    else Nothing-      Right x -> if i `Map.member` xs-                    -- if it is a pattern variable under investigation-                    -- and the e-class id is part of the trie-                    then if xs Map.! i `Set.member` _keys trie-                            -- match the next id with the next trie-                            then intersectTries var xs (_trie trie IntMap.! (xs Map.! i)) ids-                            else Nothing+      Left x  -> case IntMap.lookup x (_trie trie) of+                   Just subtrie -> intersectTries var xs subtrie ids+                   Nothing -> Nothing+      Right x -> if IntMap.member x xs+                    then case IntMap.lookup (xs IntMap.! x) (_trie trie) of+                           Just subtrie -> intersectTries var xs subtrie ids+                           Nothing -> Nothing                     else if Right x == var-                            -- not under investigation and is the var of interest                             then if all (isDiffFrom x) ids-                                    -- if there are no other occurrence of x in the next vars,-                                    -- the keys of the trie are all possible candidates-                                    then Just $ _keys trie-                                    -- oterwise, put i under investigation and check the next occurrences-                                    -- returning the intersection+                                    then Just $ Set.fromList (IntMap.keys (_trie trie))                                     else Just $ IntMap.foldrWithKey (\k v acc ->-                                                    case intersectTries var (Map.insert i k xs) v ids of+                                                    case intersectTries var (IntMap.insert x k xs) v ids of                                                       Nothing -> acc                                                       _       -> Set.insert k acc) Set.empty (_trie trie)-                            -- if it is not the var of interest-                            -- assign and test all possible e-class ids to it-                            -- and move forward                             else Just $ IntMap.foldrWithKey (\k v acc ->-                                                case intersectTries var (Map.insert i k xs) v ids of+                                                case intersectTries var (IntMap.insert x k xs) v ids of                                                   Nothing -> acc                                                   Just s  -> Set.union acc s                                                      ) Set.empty (_trie trie)@@ -359,15 +668,28 @@ {-# INLINE elemOfAtom #-}  -- | sorts the variables in a query by the most frequently occurring+-- Ties are broken by putting an atom ROOT first. The root indexes the+-- operator trie directly, so matching it first replaces repeated whole-trie+-- folds (O(candidates x nodes)) with direct per-node trie descents. The old+-- tie-break (by id) put low-id pattern leaves before the high-id fresh root,+-- which made the root's domain include every operator node regardless of the+-- already-bound children (over-enumeration and O(n^2) folds).+-- Measured on the user config: 33s -> 19s (MT -N8), best loss unchanged. orderedVars :: Query -> [ClassOrVar]-orderedVars atoms = sortBy (comparing varCost) $ nub [a | atom <- atoms, a <- getIdsFrom atom, isRight a]+orderedVars atoms = sortBy (comparing key) $ RangeSet.toList $ RangeSet.fromList [a | atom <- atoms, a <- getIdsFrom atom, isRight a]   where     getIdsFrom (Atom r t) = r : getElems t     isRight (Right _) = True     isRight _ = False +    -- is the variable the ROOT of some atom (an index into the operator trie)?+    isHeader v = any (\a -> case a of Atom r _ -> r == v) atoms+     varCost :: ClassOrVar -> Int     varCost var = foldr (\a acc -> if elemOfAtom var a then acc - 100 + atomLen a else acc) 0 atoms++    key :: ClassOrVar -> (Int, Int)+    key v = (varCost v, if isHeader v then 0 else 1)      atomLen (Atom _ t) = 1 + length (getElems t) {-# INLINE orderedVars #-}

@@ -1,9 +1,10 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE StrictData #-}-{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}+{-# LANGUAGE UndecidableInstances #-} ----------------------------------------------------------------------------- -- | -- Module      :  Algorithm.EqSat.Egraph@@ -21,115 +22,100 @@ module Algorithm.EqSat.Egraph where  import Control.Lens (element, makeLenses, view, over, (&), (+~), (-~), (.~), (^.))---import Control.Monad (forM, forM_, when, foldM, void)-import Data.List ( intercalate )+--import Control.Monad (forM_, when, foldM, void)+import Data.List ( intercalate, foldl' )+import Control.Monad (forM) import Control.Monad.State.Strict hiding ( get, put )+import Control.Monad.IO.Class (MonadIO(..))+import Data.Functor.Identity (Identity)+import GHC.Stack (HasCallStack)+import System.Random (StdGen) import Data.IntMap.Strict (IntMap) import qualified Data.IntMap.Strict as IntMap import Data.Map.Strict (Map) import qualified Data.Map.Strict as Map+import Data.HashMap.Strict (HashMap)+import qualified Data.HashMap.Strict as HashMap import Data.HashSet (HashSet) import qualified Data.HashSet as Set import Data.IntSet (IntSet) import qualified Data.IntSet as IntSet-import Data.Sequence ( Seq(..), (><) )-import qualified Data.Sequence as FingerTree-import Data.Foldable ( toList )+import qualified Data.Set as RangeSet import Data.SRTree import Data.SRTree.Eval+import Data.SRTree.Recursion (cata) import Data.Hashable import Data.Binary import qualified Data.Binary as Bin-import qualified Data.Massiv.Array as MA+import qualified Data.Vector.Unboxed as VU+import Control.DeepSeq (NFData)  import GHC.Generics -import Debug.Trace  type EClassId     = Int -- NOTE: DO NOT CHANGE THIS, this will break the use of IntMap and IntSet type ClassIdMap   = IntMap-type ENode        = SRTree EClassId-type ENodeEnc     = (Int, Int, Int, Double)++-- | N-ary operators represented as flattened multisets inside the e-graph.+-- Only Add and Mul are associative-commutative in this library; the remaining+-- ops (Sub, Div, Power, PowerAbs, AQ) stay binary and live in 'EBin'.+data NOp = EAdd | EMul deriving (Show, Eq, Ord, Enum, Generic, NFData)++-- | The e-graph's node language.+--+-- 'ENAry' stores Add/Mul as a canonical multiset of e-class ids: children are+-- path-compressed, keys sorted by canonical 'EClassId' (commutativity), and+-- nested same-op ENAry children are absorbed at insertion time+-- (associativity), so no commutativity/associativity rewrite rules are needed+-- for Add/Mul. The children are an 'IntMap' of e-class id to multiplicity.+data ENode+  = EVar   {-# UNPACK #-} !Int+  | EParam {-# UNPACK #-} !Int+  | EConst {-# UNPACK #-} !Double+  | EUni   Function EClassId+  | EBin   Op EClassId EClassId          -- Sub | Div | Power | PowerAbs | AQ+  | ENAry  NOp (IntMap Int)              -- canonical multiset: eclass -> multiplicity+  deriving (Show, Eq, Generic, NFData)+ type EGraphST m a = StateT EGraph m a type Cost         = Int type CostFun      = SRTree Cost -> Cost-type ECache = IntMap.IntMap PVector+type ECache = IntMap.IntMap Target +instance Hashable NOp where+  hashWithSalt n EAdd = n `hashWithSalt` (0 :: Int)+  hashWithSalt n EMul = n `hashWithSalt` (1 :: Int)+ instance Hashable ENode where-  hashWithSalt n enode = hashWithSalt n (encodeEnode enode)+  hashWithSalt n (EVar ix)      = n `hashWithSalt` (0 :: Int) `hashWithSalt` ix+  hashWithSalt n (EParam ix)    = n `hashWithSalt` (1 :: Int) `hashWithSalt` ix+  hashWithSalt n (EConst x)     = n `hashWithSalt` (2 :: Int) `hashWithSalt` x+  hashWithSalt n (EUni f t)     = n `hashWithSalt` (3 :: Int) `hashWithSalt` (fromEnum f) `hashWithSalt` t+  hashWithSalt n (EBin op l r)  = n `hashWithSalt` (4 :: Int) `hashWithSalt` (fromEnum op) `hashWithSalt` l `hashWithSalt` r+  hashWithSalt n (ENAry op m)   = n `hashWithSalt` (5 :: Int) `hashWithSalt` op `hashWithSalt` m -type RangeTree a = Seq (a, EClassId)+type RangeTree a = RangeSet.Set (a, EClassId) --- | this assumes up to 999 variables and params-encodeEnode :: ENode -> ENodeEnc---encodeEnode = id-{--}-encodeEnode (Var ix)         = (0, ix, -1, 0)-encodeEnode (Param ix)       = (1, ix, -1, 0)-encodeEnode (Const x)        = (2, -1, -1, x)-encodeEnode (Uni f ed)       = (300 + fromEnum f, ed, -1, 0)-encodeEnode (Bin op ed1 ed2) = (400 + fromEnum op, ed1, ed2, 0)-{--}-{-# INLINE encodeEnode #-}+-- | Expand a canonical multiset back to the equivalent (multi-)set of child+-- e-class ids, one entry per occurrence.+expandedList :: IntMap Int -> [EClassId]+expandedList = concatMap (\(k, n) -> replicate n k) . IntMap.toAscList+{-# INLINE expandedList #-} -decodeEnode :: ENodeEnc -> ENode---decodeEnode = id-{--}-decodeEnode (0, ix, _, _) = Var ix-decodeEnode (1, ix, _, _) = Param ix-decodeEnode (2, _, _, x)  = Const x-decodeEnode (opCode, arg1, arg2, arg3)-  | opCode < 400 = Uni (toEnum $ opCode-300) arg1-  | otherwise    = Bin (toEnum $ opCode-400) arg1 arg2-  {--}-{-# INLINE decodeEnode #-}+-- | Build a canonical multiset from a list of child ids (duplicates allowed).+imFromList :: [EClassId] -> IntMap Int+imFromList = IntMap.fromListWith (+) . map (, 1)+{-# INLINE imFromList #-} ++ insertRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a-insertRange eid x Empty                      = FingerTree.singleton (x, eid)-insertRange eid x (y :<| _xs) | (x, eid) < y = (x, eid) :<| y :<| _xs-insertRange eid x (_xs :|> y) | (x, eid) > y = _xs :|> y :|> (x, eid)-insertRange eid x rt = go rt-  where-    entry   = (x, eid)-    go root = case FingerTree.splitAt (n `div` 2) root of-                (Empty, Empty)    -> FingerTree.singleton entry-                (Empty, z :<| zs) | entry < z -> entry :<| z :<| zs-                                  | otherwise -> z :<| (go zs)-                (ys :|> y, Empty) | entry > y -> ys :|> y :|> entry-                                  | otherwise -> (go ys) :|> y-                (ys :|> y, z :<| zs)-                     | entry > y && entry < z -> (ys :|> y :|> entry) >< (z :<| zs)-                     | entry > z              -> (ys :|> y) >< go (z :<| zs)-                     | entry < y              -> go (ys :|> y) >< (z :<| zs)-                     | otherwise              -> root-      where-        n = FingerTree.length root+insertRange eid x = RangeSet.insert (x, eid)+{-# INLINE insertRange #-}  removeRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a-removeRange eid x Empty                  = Empty-removeRange eid x (y :<| _xs) | (x, eid) < y = (y :<| _xs)-removeRange eid x (_xs :|> y) | (x, eid) > y = (_xs :|> y)-removeRange eid x rt = go rt-  where-    entry   = (x, eid)-    go root = case FingerTree.splitAt (n `div` 2) root of-                (Empty, Empty)    -> root-                (Empty, z :<| zs)-                            | entry < z  -> z :<| zs-                            | entry == z -> zs-                            | otherwise  -> z :<| (go zs)-                (ys :|> y, Empty)-                            | entry > y  -> ys :|> y-                            | entry == y -> ys-                            | otherwise  -> (go ys) :|> y-                (ys :|> y, z :<| zs)-                     | entry > y && entry < z -> root-                     | entry > z              -> (ys :|> y) >< go (z :<| zs)-                     | entry < y              -> go (ys :|> y) >< (z :<| zs)-                     | otherwise              -> root--      where-        n = FingerTree.length root+removeRange eid x = RangeSet.delete (x, eid)+{-# INLINE removeRange #-}   @@ -137,46 +123,50 @@  -- TODO: check this \/ getWithinRange :: Ord a => a -> a -> RangeTree a -> [EClassId]-getWithinRange lb ub rt = map snd . toList $ go rt-  where-    go Empty = Empty-    go root = case FingerTree.splitAt (n `div` 2) root of-                (Empty, Empty)    -> Empty-                (ys :|> y, Empty)-                     | fst y < lb    -> Empty-                     | otherwise -> go (ys :|> y)-                (Empty, z :<| zs)-                            | fst z > ub    -> Empty-                            | otherwise -> go (z :<| zs)-                (ys :|> y, z :<| zs)-                     | fst y < lb -> go (z :<| zs)-                     | fst z > ub -> go (ys :|> y)-                     | otherwise -> go (ys :|> y) >< go (z :<| zs)-      where-        n = FingerTree.length root-+getWithinRange lb ub rt =+  let (_, ge)  = RangeSet.split (lb, minBound) rt+      (inR, _) = RangeSet.split (ub, maxBound) ge+  in map snd (RangeSet.toList inR) -getSmallest :: Ord a => RangeTree a -> (a, EClassId)-getSmallest rt = case rt of-                     Empty -> error "empty finger"-                     x :<| t -> x+getSmallest :: Ord a => RangeTree a -> Maybe (a, EClassId)+getSmallest = RangeSet.lookupMin {-# INLINE getSmallest #-} -getGreatest :: Ord a => RangeTree a -> (a, EClassId)-getGreatest rt = case rt of-                     Empty -> error "empty finger"-                     t :|> x -> x+getGreatest :: Ord a => RangeTree a -> Maybe (a, EClassId)+getGreatest = RangeSet.lookupMax {-# INLINE getGreatest #-} +-- | Handle to an external, lazily paged e-class store (provided by the+-- storage layer, e.g. srtree-db's 'PageStore'). An 'EGraph' carries one when+-- e-classes are backed by a database; the IO actions fetch / persist /+-- evict a single e-class page. 'Nothing' keeps the classic fully-resident+-- behaviour.+data EClassPageStore = EClassPageStore+  { cpsLookup :: EClassId -> IO (Maybe EClass)+  , cpsInsert :: EClass -> IO ()+  , cpsDelete :: EClassId -> IO ()+  , cpsFlush  :: IO ()                      -- ^ write back all pending dirty pages+  , cpsAll    :: IO [EClass]                -- ^ all e-classes currently in the store+  , cpsKeys   :: IO [EClassId]              -- ^ all e-class ids currently in the store+  , cpsStreamRoots :: SRTree () -> Int -> [EClassId] -> IO [EClassId]  -- ^ bounded candidate roots for an operator, skipping an attempted set+  , cpsRecordNode  :: ENode -> EClassId -> IO ()         -- ^ register a newly-created node for write-back+  , cpsNodeToClass :: ENode -> IO (Maybe EClassId)       -- ^ content-address node -> class lookup (live)+  , cpsCanonicalOf :: EClassId -> IO (Maybe EClassId)    -- ^ e-class -> canonical representative (live)+  , cpsRecordCanonical :: EClassId -> EClassId -> IO ()  -- ^ persist a canonical mapping (write-back)+  , cpsBeginFrontier :: IO ()                            -- ^ start a frontier re-saturation (restrict matcher to changed classes)+  , cpsEndFrontier    :: IO ()                           -- ^ end it: clear the frontier (a pass re-saturated everything)+  }+ data EGraph = EGraph { _canonicalMap  :: ClassIdMap EClassId   -- maps an e-class id to its canonical form-                     , _eNodeToEClass :: Map ENode EClassId    -- maps an e-node to its e-class id-                     , _eClass        :: ClassIdMap EClass     -- maps an e-class id to its e-class data+                     , _eNodeToEClass :: HashMap ENode EClassId    -- maps an e-node to its e-class id+                     , _eClass        :: ClassIdMap EClass     -- maps an e-class id to its e-class data (resident cache)                      , _eDB           :: EGraphDB-                     } deriving (Show, Generic)+                     , _classStore    :: Maybe EClassPageStore -- optional lazily paged store for _eClass+                     }  data EGraphDB = EDB { _worklist      :: HashSet (EClassId, ENode)      -- e-nodes and e-class schedule for analysis                     , _analysis      :: HashSet (EClassId, ENode)      -- e-nodes and e-class that changed data-                    , _refits        :: HashSet EClassId+                     , _refits        :: IntSet                     , _patDB         :: DB                         -- database of patterns                     , _fitRangeDB    :: RangeTree Double           -- database of valid fitness                     , _dlRangeDB     :: RangeTree Double@@ -184,26 +174,29 @@                     , _sizeFitDB     :: IntMap (RangeTree Double)  -- hacky! Size x Fitness DB                     , _sizeDLDB      :: IntMap (RangeTree Double)                     , _unevaluated   :: IntSet                     -- set of not-evaluated e-classes-                    , _nextId        :: Int                        -- next available id-                    } deriving (Show, Generic)+                      , _nextId        :: Int                        -- next available id+                      , _changed       :: !Bool                      -- dirty flag: true if modified since last check+                      , _trackDBs      :: !Bool                      -- maintain range DBs (False during pure simplify)+                      , _seenMatches   :: Map String (RangeSet.Set String) -- persistent (rule source -> attempted match keys)+                      } deriving (Show, Generic) -data EClass = EClass { _eClassId :: Int                   -- e-class id (maybe we don't need that here)-                     , _eNodes   :: HashSet ENodeEnc          -- set of e-nodes inside this e-class+data EClass = EClass { _eClassId :: {-# UNPACK #-} !Int                   -- e-class id (maybe we don't need that here)+                     , _eNodes   :: HashSet ENode           -- set of e-nodes inside this e-class                      , _parents  :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s-                     , _height   :: Int                   -- height+                     , _height   :: {-# UNPACK #-} !Int                   -- height                      , _info     :: EClassData            -- data                      } deriving (Show, Eq, Generic) -data Consts   = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq, Generic)+data Consts   = NotConst | ParamIx {-# UNPACK #-} !Int | ConstVal {-# UNPACK #-} !Double deriving (Show, Eq, Generic) data Property = Positive | Negative | NonZero | Real deriving (Show, Eq, Generic) -- TODO: incorporate properties -data EClassData = EData { _cost    :: Cost+data EClassData = EData { _cost    :: {-# UNPACK #-} !Cost                         , _best    :: ENode                         , _consts  :: Consts                         , _fitness :: Maybe Double    -- NOTE: this cannot be NaN                         , _dl      :: Maybe Double-                        , _theta   :: [PVector]-                        , _size    :: Int+                        , _theta   :: [Target]+                        , _size    :: {-# UNPACK #-} !Int                         -- , _properties :: Property                         -- TODO: include evaluation of expression from this e-class                         } deriving (Show, Generic)@@ -211,21 +204,32 @@ -- * Serialization instance Generic (EClassId, ENode) -instance Binary (SRTree EClassId) where-  put (Var ix)     = put (0 :: Word8) >> put ix-  put (Param ix)   = put (1 :: Word8) >> put ix-  put (Const x)    = put (2 :: Word8) >> put x-  put (Uni f t)    = put (3 :: Word8) >> put (fromEnum f) >> put t-  put (Bin op l r) = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r+instance Binary NOp where+  put EAdd = put (0 :: Word8)+  put EMul = put (1 :: Word8)    get = do t <- get :: Get Word8            case t of-                0 -> Var   <$> get-                1 -> Param <$> get-                2 -> Const <$> get-                3 -> Uni   <$> (toEnum <$> get) <*> get-                4 -> Bin   <$> (toEnum <$> get) <*> get <*> get+             0 -> pure EAdd+             1 -> pure EMul +instance Binary ENode where+  put (EVar ix)      = put (0 :: Word8) >> put ix+  put (EParam ix)    = put (1 :: Word8) >> put ix+  put (EConst x)     = put (2 :: Word8) >> put x+  put (EUni f t)     = put (3 :: Word8) >> put (fromEnum f) >> put t+  put (EBin op l r)  = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r+  put (ENAry op m)   = put (5 :: Word8) >> put op >> put (expandedList m)++  get = do t <- get :: Get Word8+           case t of+                0 -> EVar   <$> get+                1 -> EParam <$> get+                2 -> EConst <$> get+                3 -> EUni   <$> (toEnum <$> get) <*> get+                4 -> EBin   <$> (toEnum <$> get) <*> get <*> get+                5 -> ENAry  <$> get <*> (imFromList <$> get)+ instance Binary (SRTree ()) where   put (Var ix)     = put (0 :: Word8) >> put ix   put (Param ix)   = put (1 :: Word8) >> put ix@@ -245,17 +249,30 @@   put hs = put (Set.toList hs)   get    = Set.fromList <$> get -instance Binary PVector where-  put xs = put (MA.toList xs)-  get    = MA.fromList compMode <$> get+instance (Binary k, Binary v, Hashable k, Eq k) => Binary (HashMap k v) where+  put hm = put (HashMap.toList hm)+  get    = HashMap.fromList <$> get +instance Binary Target where+  put xs = put (VU.toList xs)+  get    = VU.fromList <$> get+ instance Binary IntTrie instance Binary EClass instance Binary Consts instance Binary Property instance Binary EClassData-instance Binary EGraphDB-instance Binary EGraph+-- Custom: keep `_trackDBs` out of the wire format so on-disk EGraphDB data+-- (written before the flag existed) decodes unchanged; it defaults to True.+instance Binary EGraphDB where+  put (EDB w a r p f d s sf sdl u n c _ _) =+    put w >> put a >> put r >> put p >> put f >> put d >> put s >> put sf >> put sdl >> put u >> put n >> put c+  get = EDB <$> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> pure True <*> pure Map.empty+-- Custom: the wire format omits `_classStore` (a runtime handle to the paged+-- store, never serialized); it decodes to Nothing.+instance Binary EGraph where+  put (EGraph c n e d _) = put c >> put n >> put e >> put d+  get = EGraph <$> get <*> get <*> get <*> get <*> pure Nothing  instance Eq EClassData where   EData c1 b1 cs1 ft1 dl1 _ s1 == EData c2 b2 cs2 ft2 dl2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && dl1==dl2 && s1==s2@@ -267,92 +284,642 @@ -- The IntTrie is composed of the set of available keys (for convenience) -- and an IntMap that maps one e-class id to the first child IntTrie, -- the first child IntTrie will point to the next child and so on-data IntTrie = IntTrie { _keys :: HashSet EClassId, _trie :: IntMap IntTrie } deriving (Generic)+newtype IntTrie = IntTrie { _trie :: IntMap IntTrie } deriving (Generic) --- Shows the IntTrie as {keys} -> {show IntTries} instance Show IntTrie where-  show (IntTrie k t) = let keys  = intercalate "," (map show $ Set.toList k)-                           tries = intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t)-                       in "{" <> keys <> "} - {" <> tries <> "}"+  show (IntTrie t) = "{" <> intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t) <> "}"  makeLenses ''EGraph makeLenses ''EClass makeLenses ''EClassData makeLenses ''EGraphDB +-- * Paged e-class access++-- | A monad that can serve e-class data.+--+-- The pure instances ('Identity', 'State StdGen') serve classes from the+-- resident @_eClass@ map; the 'MonadIO' instance consults the optional+-- 'EClassPageStore' when the graph carries one, falling back to the resident+-- map otherwise. All e-class read/write goes through these accessors, which+-- are the single choke point for a paged (out-of-core) e-graph.+class Monad m => ClassStore m where+  lookupClass :: EClassId -> EGraphST m (Maybe EClass)+  getClass    :: HasCallStack => EClassId -> EGraphST m EClass+  insertClass :: EClass -> EGraphST m ()+  deleteClass :: EClassId -> EGraphST m ()+  adjustClass :: EClassId -> (EClass -> EClass) -> EGraphST m ()+  -- | Enumerate every e-class (ids / values) in the graph. Paged graphs stream+  -- from the store; resident graphs read the full @_eClass@ map.+  allClasses  :: EGraphST m [EClass]+  allKeys     :: EGraphST m [EClassId]+  -- | Read/write a class directly from/to the backing store, bypassing the+  -- resident LRU cache (and its O(n) 'trimResidentCache'). Bulk single-pass+  -- traversals such as 'recalculateBestAllStream' must use these: routing every+  -- one of ~n classes through 'lookupClass'/'insertClass' inserts each into the+  -- resident map and calls 'trimResidentCache' (a full O(n) rebuild) after each+  -- write, degenerating to O(n^2) and never terminating at scale.+  readDirect  :: EClassId -> EGraphST m (Maybe EClass)+  writeDirect :: EClass -> EGraphST m ()+  allClasses  = gets (IntMap.elems . _eClass)+  allKeys     = gets (IntMap.keys . _eClass)+  readDirect  = lookupClass+  writeDirect = insertClass+  -- | Enumerate (bounded) candidate e-class ids that contain a node of the+  -- given operator, to drive the streaming matcher, skipping any ids in+  -- @exclude@ (the already-attempted seen-set, so the per-rule budget advances+  -- to new roots across scheduler cycles). The default reads the resident+  -- @_patDB@ trie (the fully-in-RAM path); a paged graph streams the candidates+  -- from its backing store instead, so the matcher never builds an O(nodes)+  -- structure.+  streamRoots :: SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId]+  streamRoots = streamRootsFromDB+  -- | Record a newly-created e-node (and its e-class) so a streaming matcher's+  -- candidate source can see it. The default (fully resident graph) is a no-op:+  -- the resident @_patDB@ is already updated by 'addToDB'.+  recordNode :: ENode -> EClassId -> EGraphST m ()+  recordNode _ _ = pure ()+  -- | Content-address node -> class lookup. The default reads the resident+  -- @_eNodeToEClass@ map (complete for a resident graph); a paged graph bounds+  -- that map and falls back to the backing store on a miss.+  lookupNode :: ENode -> EGraphST m (Maybe EClassId)+  lookupNode en = gets (HashMap.lookup en . _eNodeToEClass)+  -- | Record a node -> class mapping. The default keeps the resident (full)+  -- map; a paged graph bounds it (evicting, since the store is authoritative).+  insertNode :: ENode -> EClassId -> EGraphST m ()+  insertNode en eid = modify' $ over eNodeToEClass (HashMap.insert en eid)+  -- | Record a canonical mapping (e-class -> representative), persisting it on a+  -- paged graph so the store-backed canonical lookup sees merges/new classes.+  insertCanonical :: EClassId -> EClassId -> EGraphST m ()+  insertCanonical eid canon = modify' $ over canonicalMap (IntMap.insert eid canon)+  -- | The canonical representative of an e-class, or @Nothing@ when unknown. The+  -- default reads the resident @_canonicalMap@; a paged graph bounds it and+  -- falls back to the store.+  canonicalOf :: EClassId -> EGraphST m (Maybe EClassId)+  canonicalOf eid = gets (IntMap.lookup eid . _canonicalMap)++-- | Default candidate-root enumeration from the resident @_patDB@ trie, capped+-- at @budget@ after skipping @exclude@ (used by the pure instances and as the+-- no-store fallback for a @MonadIO@ graph).+streamRootsFromDB :: Monad m => SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId]+streamRootsFromDB op budget exclude = do+  db <- gets (_patDB . _eDB)+  let ex = IntSet.fromList exclude+  case Map.lookup op db of+    Nothing  -> pure []+    Just trie -> pure (take budget [ e | e <- IntMap.keys (_trie trie), not (IntSet.member e ex) ])+{-# INLINE streamRootsFromDB #-}++-- | Whether the graph is backed by a lazily paged e-class store. Streaming+-- matchers dispatch on this: a paged graph enumerates candidates from the+-- backing store (bounded memory), a resident graph from @_patDB@.+isPagedGraph :: Monad m => EGraphST m Bool+isPagedGraph = gets (maybe False (const True) . _classStore)+{-# INLINE isPagedGraph #-}++-- Resident-map implementations (used by every pure monad) ------------------++pureLookupClass :: Monad m => EClassId -> EGraphST m (Maybe EClass)+pureLookupClass cid = gets (IntMap.lookup cid . _eClass)+{-# INLINE pureLookupClass #-}++pureGetClass :: (Monad m, HasCallStack) => EClassId -> EGraphST m EClass+pureGetClass cid = do+  m <- pureLookupClass cid+  case m of+    Just ec -> pure ec+    Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid+{-# INLINE pureGetClass #-}++pureInsertClass :: Monad m => EClass -> EGraphST m ()+pureInsertClass ec = modify' $ over eClass (IntMap.insert (_eClassId ec) ec)+{-# INLINE pureInsertClass #-}++pureDeleteClass :: Monad m => EClassId -> EGraphST m ()+pureDeleteClass cid = modify' $ over eClass (IntMap.delete cid)+{-# INLINE pureDeleteClass #-}++pureAdjustClass :: Monad m => EClassId -> (EClass -> EClass) -> EGraphST m ()+pureAdjustClass cid f = modify' $ over eClass (IntMap.adjust f cid)+{-# INLINE pureAdjustClass #-}++-- | Maximum number of e-classes kept in the resident @_eClass@ cache when the+-- graph is backed by a paged store. When exceeded, the largest-id classes are+-- retained and the rest evicted from the resident map. The store remains+-- authoritative (and Little-data reads fall back to it), so eviction only+-- bounds memory, never correctness.+residentClassCap :: Int+residentClassCap = 50000++-- | Trim the resident @_eClass@ cache to at most 'residentClassCap' entries+-- by keeping the largest ids. No-op for graphs without a paged store (their+-- resident map must stay complete for the pure instances). Halving on 2x keeps+-- steady churn from triggering an O(n) rebuild on every insert.+trimResidentCache :: Monad m => EGraphST m ()+trimResidentCache = modify' $ \eg ->+  case _classStore eg of+    Nothing -> eg+    Just _  ->+      let m = _eClass eg+          n = IntMap.size m+      in if n <= 2 * residentClassCap+            then eg+            else over eClass (const (IntMap.fromList (Prelude.drop (n - residentClassCap) (IntMap.toAscList m)))) eg++-- | Bound on the resident @_eNodeToEClass@ cache on a paged graph. Beyond the+-- cap (checked at 2x, halved back to cap) the map is pruned; the backing store+-- is authoritative, so eviction only trades a little dedup accuracy for bounded+-- memory, never correctness.+nodeCacheCap :: Int+nodeCacheCap = 100000++-- | Bound on the resident @_canonicalMap@ cache on a paged graph (same+-- halve-on-2x policy; evicted entries are re-read from the store).+canonicalCacheCap :: Int+canonicalCacheCap = 100000+{-# INLINE nodeCacheCap #-}+{-# INLINE canonicalCacheCap #-}++trimNodeCache :: Monad m => EGraphST m ()+trimNodeCache = modify' $ \eg ->+  case _classStore eg of+    Nothing -> eg+    Just _  ->+      let m = _eNodeToEClass eg+          n = HashMap.size m+      in if n <= 2 * nodeCacheCap+            then eg+            else over eNodeToEClass (const (HashMap.fromList (Prelude.take nodeCacheCap (HashMap.toList m)))) eg+{-# INLINE trimNodeCache #-}++trimCanonicalCache :: Monad m => EGraphST m ()+trimCanonicalCache = modify' $ \eg ->+  case _classStore eg of+    Nothing -> eg+    Just _  ->+      let m = _canonicalMap eg+          n = IntMap.size m+      in if n <= 2 * canonicalCacheCap+            then eg+            else over canonicalMap (const (IntMap.fromList (Prelude.take canonicalCacheCap (IntMap.toAscList m)))) eg+{-# INLINE trimCanonicalCache #-}++instance ClassStore Identity where+  lookupClass = pureLookupClass+  getClass    = pureGetClass+  insertClass = pureInsertClass+  deleteClass = pureDeleteClass+  adjustClass = pureAdjustClass++instance ClassStore (State StdGen) where+  lookupClass = pureLookupClass+  getClass    = pureGetClass+  insertClass = pureInsertClass+  deleteClass = pureDeleteClass+  adjustClass = pureAdjustClass++-- Any monad that can run IO is potentially paged: the graph's optional+-- store, when present, is authoritative; otherwise classes come from the+-- resident map.+instance {-# OVERLAPPABLE #-} (Monad m, MonadIO m) => ClassStore m where+  -- The resident map is kept in sync by 'insertClass'/'deleteClass', so it is+  -- consulted first: repeated reads never touch the store, and a class that+  -- was evicted from the store's LRU while still dirty is never served stale.+  lookupClass cid = do+    eg <- gets id+    case IntMap.lookup cid (_eClass eg) of+      Just ec -> pure (Just ec)+      Nothing -> case _classStore eg of+                   Nothing -> pure Nothing+                   Just h  -> liftIO (cpsLookup h cid)+  getClass cid = do+    eg <- gets id+    case IntMap.lookup cid (_eClass eg) of+      Just ec -> pure ec+      Nothing -> case _classStore eg of+                   Nothing -> pureGetClass cid+                   Just h  -> do+                     m <- liftIO (cpsLookup h cid)+                     case m of+                       Just ec -> do+                         modify' (over eClass (IntMap.insert cid ec))+                         trimResidentCache+                         pure ec+                       Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid+  insertClass ec = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pureInsertClass ec+      Just h  -> do liftIO (cpsInsert h ec)+                    pureInsertClass ec+                    trimResidentCache+  deleteClass cid = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pureDeleteClass cid+      Just h  -> do liftIO (cpsDelete h cid)+                    pureDeleteClass cid+  adjustClass cid f = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pureAdjustClass cid f+      Just _  -> do+        m <- lookupClass cid+        case m of+          Nothing -> pure ()+          Just ec -> insertClass (f ec)+  allClasses = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pure (IntMap.elems (_eClass eg))+      Just h  -> liftIO (cpsAll h)+  allKeys = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pure (IntMap.keys (_eClass eg))+      Just h  -> liftIO (cpsKeys h)+  -- Bypass the resident cache entirely: read the page straight from the store+  -- and never insert into the (bounded) resident map, so a bulk traversal over+  -- every class stays O(n) instead of O(n^2).+  readDirect cid = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pureLookupClass cid+      Just h  -> liftIO (cpsLookup h cid)+  writeDirect ec = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pureInsertClass ec+      Just h  -> liftIO (cpsInsert h ec)+  streamRoots op budget exclude = do+    eg <- gets id+    case _classStore eg of+      Nothing -> streamRootsFromDB op budget exclude+      Just h  -> liftIO (cpsStreamRoots h op budget exclude)+  recordNode en eid = do+    eg <- gets id+    case _classStore eg of+      Nothing -> pure ()+      Just h  -> liftIO (cpsRecordNode h en eid)+  lookupNode en = do+    eg <- gets id+    case _classStore eg of+      Nothing -> gets (HashMap.lookup en . _eNodeToEClass)+      Just h  -> do+        m <- gets (HashMap.lookup en . _eNodeToEClass)+        case m of+          Just eid -> pure (Just eid)+          Nothing -> do+            r <- liftIO (cpsNodeToClass h en)+            case r of+              Just eid -> do insertNode en eid+                             pure (Just eid)+              Nothing  -> pure Nothing+  insertNode en eid = do+    eg <- gets id+    case _classStore eg of+      Nothing -> modify' $ over eNodeToEClass (HashMap.insert en eid)+      Just _  -> do modify' $ over eNodeToEClass (HashMap.insert en eid)+                    trimNodeCache+  insertCanonical eid canon = do+    eg <- gets id+    case _classStore eg of+      Nothing -> modify' $ over canonicalMap (IntMap.insert eid canon)+      Just h  -> do modify' $ over canonicalMap (IntMap.insert eid canon)+                    trimCanonicalCache+                    liftIO (cpsRecordCanonical h eid canon)+  canonicalOf eid = do+    eg <- gets id+    case _classStore eg of+      Nothing -> gets (IntMap.lookup eid . _canonicalMap)+      Just h  -> do+        m <- gets (IntMap.lookup eid . _canonicalMap)+        case m of+          Just c  -> pure (Just c)+          Nothing -> do+            r <- liftIO (cpsCanonicalOf h eid)+            case r of+              Just c  -> do modify' $ over canonicalMap (IntMap.insert eid c)+                            trimCanonicalCache+                            pure (Just c)+              Nothing -> pure Nothing+ -- * E-Graph basic supporting functions  -- | returns an empty e-graph emptyGraph :: EGraph-emptyGraph = EGraph IntMap.empty Map.empty IntMap.empty emptyDB+emptyGraph = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDB Nothing {-# INLINE emptyGraph #-}  -- | returns an empty e-graph DB emptyDB :: EGraphDB-emptyDB = EDB Set.empty Set.empty Set.empty Map.empty FingerTree.empty FingerTree.empty IntMap.empty IntMap.empty IntMap.empty IntSet.empty 0+emptyDB = EDB+  Set.empty+  Set.empty+  IntSet.empty+  Map.empty+  RangeSet.empty+  RangeSet.empty+  IntMap.empty+  IntMap.empty+  IntMap.empty+  IntSet.empty+  0+  False+  True+  Map.empty {-# INLINE emptyDB #-} +-- | like 'emptyDB' but skips range-DB maintenance (pure simplify mode)+emptyDBNoTrack :: EGraphDB+emptyDBNoTrack = emptyDB{ _trackDBs = False }+{-# INLINE emptyDBNoTrack #-}++-- | an empty e-graph that skips range-DB maintenance (pure simplify mode)+emptyGraphNoTrack :: EGraph+emptyGraphNoTrack = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDBNoTrack Nothing+{-# INLINE emptyGraphNoTrack #-}+ -- | Creates a new e-class from an e-class id, a new e-node, -- and the info of this e-class  createEClass :: EClassId -> ENode -> EClassData -> Int -> EClass-createEClass cId enode' info h = EClass cId (Set.singleton $ encodeEnode enode') Set.empty h info+createEClass cId enode' info h = EClass cId (Set.singleton enode') Set.empty h info {-# INLINE createEClass #-} --- | gets the canonical id of an e-class-canonical :: Monad m => EClassId -> EGraphST m EClassId-canonical eclassId =-  do m <- gets _canonicalMap-     let oneStep = m IntMap.! eclassId-     if oneStep == eclassId-        then pure eclassId-        else go m oneStep-    where-      go :: Monad m => IntMap EClassId -> EClassId -> EGraphST m EClassId-      go m ecId-        | m IntMap.! ecId == ecId = do modify' $ over canonicalMap (IntMap.insert eclassId ecId) -- creates a shortcut for next time-                                       pure ecId        -- if the e-class id is mapped to itself, it's canonical-        | otherwise        = go m (m IntMap.! ecId)  -- otherwise, check the next id in the sequence+-- | gets the canonical id of an e-class with full path compression+canonical :: ClassStore m => EClassId -> EGraphST m EClassId+canonical eclassId = do+  mStep <- canonicalOf eclassId+  case mStep of+    Nothing -> canonError eclassId+    Just oneStep+      | oneStep == eclassId -> pure eclassId+      | otherwise -> do+          (root, chain) <- walk [eclassId] oneStep+          -- compress the chain in the resident cache (cache-only: the store+          -- keeps the authoritative semantic mappings recorded at insert+          -- time, so eviction just loses the shortcut, never correctness).+          modify' $ \eg -> eg{ _canonicalMap =+                        foldl' (\m' k -> IntMap.insert k root m') (_canonicalMap eg) chain }+          pure root+  where+    walk :: ClassStore m => [EClassId] -> EClassId -> EGraphST m (EClassId, [EClassId])+    walk chain ecId = do+      mNext <- canonicalOf ecId+      case mNext of+        Nothing -> canonError ecId+        Just n+          | n == ecId -> pure (ecId, chain)+          | otherwise -> walk (ecId : chain) n++    canonError :: ClassStore m => EClassId -> EGraphST m a+    canonError eid = do+      m <- gets _canonicalMap+      error $ "CANON_MISSING eid=" <> show eid <> " mapSize=" <> show (IntMap.size m) {-# INLINE canonical #-}  -- | canonize the e-node children-canonize :: Monad m => ENode -> EGraphST m ENode-canonize = mapM canonical  -- applies canonical to the children+canonize :: (ClassStore m, HasCallStack) => ENode -> EGraphST m ENode+canonize (EVar ix)     = pure (EVar ix)+canonize (EParam ix)   = pure (EParam ix)+canonize (EConst x)    = pure (EConst x)+canonize (EUni f t)    = EUni f <$> canonical t+canonize (EBin op l r) = EBin op <$> canonical l <*> canonical r+-- re-map children to their canonical ids; IntMap keeps keys sorted, so+-- commutativity is structural, no rewrite rule required.+canonize (ENAry op m) = do+  m' <- IntMap.fromListWith (+) <$> forM (IntMap.toList m) (\(c, n) -> do+            c' <- canonical c+            pure (c', n))+  pure (ENAry op m') {-# INLINE canonize #-} --- | gets an e-class with id `c`-getEClass :: Monad m => EClassId -> EGraphST m EClass-getEClass c = gets ((IntMap.! c) . _eClass)+-- | The children e-class ids of an e-node.+eChildren :: ENode -> [EClassId]+eChildren (EVar _)     = []+eChildren (EParam _)   = []+eChildren (EConst _)   = []+eChildren (EUni _ t)   = [t]+eChildren (EBin _ l r) = [l, r]+eChildren (ENAry _ m)  = expandedList m+{-# INLINE eChildren #-}++toOp :: NOp -> Op+toOp EAdd = Add+toOp EMul = Mul+{-# INLINE toOp #-}++-- | Operator shape key used to index the pattern database. ENAry maps back to+-- the corresponding binary operator shape so existing (binary) Add/Mul+-- patterns address the same trie.+eOpKey :: ENode -> SRTree ()+eOpKey (EVar ix)     = Var ix+eOpKey (EParam ix)   = Param ix+eOpKey (EConst x)    = Const x+eOpKey (EUni f _)    = Uni f ()+eOpKey (EBin op _ _) = Bin op () ()+eOpKey (ENAry EAdd _) = Bin Add () ()+eOpKey (ENAry EMul _) = Bin Mul () ()+{-# INLINE eOpKey #-}++-- | Convert an e-node (children still as e-class ids) into the equivalent+-- binary SRTree shape. NOTE: only called on non-ENary nodes; flattened+-- ENAry nodes have no binary skeleton (see 'naryTree' / the explicit ENAry+-- cases in the analyses).+fromENode :: ENode -> SRTree EClassId+fromENode (EVar ix)     = Var ix+fromENode (EParam ix)   = Param ix+fromENode (EConst x)    = Const x+fromENode (EUni f t)    = Uni f t+fromENode (EBin op l r) = Bin op l r+fromENode (ENAry _ _)   = error "fromENode: ENAry has no binary skeleton"+{-# INLINE fromENode #-}++-- | Right-fold a list of e-class child expressions into a binary Fix SRTree+-- for a flattened ENAry multiset (extraction).+naryTree :: NOp -> [Fix SRTree] -> Fix SRTree+naryTree op ts = normalizeSubDiv (foldr1 (\a b -> Fix (Bin (toOp op) a b)) ts)+{-# INLINE naryTree #-}++-- | Re-render the internal negate/recip canonical forms back as Sub/Div so+-- extraction output keeps the familiar shape: `x + (-1)*y` -> `x - y`,+-- `x + (-3)` -> `x - 3` and `x * recip y` -> `x / y`. Sub and Div never+-- appear as e-nodes; they only reappear here during reconstruction.+normalizeSubDiv :: Fix SRTree -> Fix SRTree+normalizeSubDiv = cata alg+  where+    alg :: SRTree (Fix SRTree) -> Fix SRTree+    alg (Bin Add l r) = case pick l r of+        Just (pos, neg) -> Fix (Bin Sub pos neg)+        Nothing         -> Fix (Bin Add l r)+      where+        pick a b = case negated a of+                     Just t -> Just (b, t)+                     Nothing -> case negated b of+                                  Just t -> Just (a, t)+                                  Nothing -> Nothing+        negated (Fix (Bin Mul (Fix (Const c)) t)) | c == -1 = Just t+        negated (Fix (Bin Mul t (Fix (Const c)))) | c == -1 = Just t+        negated (Fix (Const c)) | c < 0 = Just (Fix (Const (-c)))+        negated _ = Nothing+    alg (Bin Mul l r) = case pick l r of+        Just (num, den) -> Fix (Bin Div num den)+        Nothing         -> Fix (Bin Mul l r)+      where+        pick a b = case a of+                     Fix (Uni Recip t) -> Just (b, t)+                     _ -> case b of+                            Fix (Uni Recip t) -> Just (a, t)+                            _ -> Nothing+    alg t = Fix t++-- | Convert a binary SRTree (children as e-class ids) into an e-node,+-- flattening Add/Mul into canonical ENAry multisets.+toENode :: (ClassStore m, HasCallStack) => SRTree EClassId -> EGraphST m ENode+toENode (Var ix)     = pure (EVar ix)+toENode (Param ix)   = pure (EParam ix)+toENode (Const x)    = pure (EConst x)+toENode (Uni f t)    = EUni f <$> canonical t+toENode (Bin Add l r) = mkENary EAdd [l, r]+toENode (Bin Mul l r) = mkENary EMul [l, r]+toENode (Bin op l r)  = EBin op <$> canonical l <*> canonical r+toENode n             = error $ "toENode: unsupported node " <> show n+{-# INLINE toENode #-}++-- | Build a canonical ENAry from child ids: canonicalize children, absorb+-- nested same-op ENAry children (associativity), sort by key (commutativity).+mkENary :: (ClassStore m, HasCallStack) => NOp -> [EClassId] -> EGraphST m ENode+mkENary op cids = mkENaryM op (imFromList cids)++-- | Build a canonical ENAry from a canonical multiset of child ids.+mkENaryM :: (ClassStore m, HasCallStack) => NOp -> IntMap Int -> EGraphST m ENode+mkENaryM op m = do+  flat <- IntMap.unionsWith (+) <$> mapM (expandM op) (IntMap.toList m)+  pure (ENAry op flat)++-- | If the e-class of `cid` holds exactly one e-node and that node is an ENAry+-- of the same op, return its children scaled by `n` (flattening `n`+-- occurrences); otherwise return `n` copies of `cid`. Flattening is only sound+-- through a class with a single node: if the class were merged with other+-- nodes (e.g. `{Add[a,b], Mul[x,c]}`) flattening would silently pick one+-- representative and change the meaning of the term.+expandM :: (ClassStore m, HasCallStack) => NOp -> (EClassId, Int) -> EGraphST m (IntMap Int)+expandM op (cid, n) = do+  ec <- getEClass cid+  case Set.toList (_eNodes ec) of+    [ENAry op' m'] | op' == op -> pure (IntMap.map (* n) m')+    _                          -> pure (IntMap.singleton cid n)++-- | Reconstruct a binary Fix SRTree from an e-node, right-folding ENAry+-- into nested Bin Add/Mul.+enodeToTree :: (ClassStore m, HasCallStack) => ENode -> EGraphST m (Fix SRTree)+enodeToTree (EVar ix)   = pure (Fix (Var ix))+enodeToTree (EParam ix) = pure (Fix (Param ix))+enodeToTree (EConst x)  = pure (Fix (Const x))+enodeToTree (EUni f t)  = Fix . Uni f <$> getBestExpr t+enodeToTree (EBin op l r) = do+  tl <- getBestExpr l+  tr <- getBestExpr r+  pure (Fix (Bin op tl tr))+enodeToTree (ENAry op m) = do+  ts <- mapM getBestExpr (expandedList m)+  pure (naryTree op ts)+{-# INLINE enodeToTree #-}++-- | gets an e-class with id `c` (auto-canonizes)+getEClass :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m EClass+getEClass c = do c' <- canonical c; getClass c' {-# INLINE getEClass #-} +-- | gets the best expression given the default cost function. Cycle-safe and+-- budgeted: see 'getBestExprBounded'.+getBestExpr :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree)+getBestExpr eid = getBestExprBounded eid++-- | Like 'getBestExpr' but terminates on pathological graphs: a visited set+-- stops the expansion from re-entering an already-expanded class (a @_best@+-- cycle arising from supersaturation/merges), and a node budget caps the total+-- expanded size (so an exponentially-shared DAG is truncated rather than+-- exploded). Both guards substitute a @Var 0@ placeholder for the part that+-- would otherwise blow up. On well-formed acyclic graphs with small bests+-- neither guard triggers, so the result is identical to the unbounded version.+-- This keeps out-of-core extraction (e.g. 'dbTop') bounded in memory.+getBestExprBounded :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree)+getBestExprBounded eid = fst <$> expand Set.empty 0 eid+  where+    budget :: Int+    budget = 200+    -- expand returns the tree and the running count of expanded nodes, so the+    -- budget bounds the TOTAL size (not just the depth): an exponentially-shared+    -- DAG is truncated instead of exploded. A revisited (cyclic) class or a+    -- full budget yields a @Var 0@ placeholder.+    expand :: ClassStore m => HashSet EClassId -> Int -> EClassId -> EGraphST m (Fix SRTree, Int)+    expand _ n _ | n >= budget = pure (Fix (Var 0), n)+    expand seen n eid+      | Set.member eid seen = pure (Fix (Var 0), n)+      | otherwise = do+          best <- (_best . _info) <$> getEClass eid+          let seen' = Set.insert eid seen+              n0    = n + 1+          case best of+            EVar ix   -> pure (Fix (Var ix), n0)+            EParam ix -> pure (Fix (Param ix), n0)+            EConst x  -> pure (Fix (Const x), n0)+            EUni f t  -> do (tt, n1) <- expand seen' n0 t+                            pure (Fix (Uni f tt), n1)+            EBin op l r -> do+              (tl, n1) <- expand seen' n0 l+              (tr, n2) <- expand seen' n1 r+              pure (Fix (Bin op tl tr), n2)+            ENAry op m -> do+              (xs, nEnd) <- goNary seen' n0 (IntMap.toAscList m) []+              pure (if null xs then (Fix (Var 0), nEnd) else (naryTree op xs, nEnd))+    -- build the ENAry children from the multiset WITHOUT materialising the+    -- expanded multiplicity list: an enormous count (a pathological supersaturated+    -- class) is capped per-child and by the total budget, so each copy counts+    -- toward the budget and no giant list is ever allocated.+    goNary seen n es acc+      | n >= budget = pure (reverse acc, n)+      | otherwise = case es of+          [] -> pure (reverse acc, n)+          ((c, cnt) : rest) -> do+            (t, n1) <- expand seen n c+            let take = min cnt (budget - n1 + 1)+                n2   = n1 + (take - 1)+                acc' = Prelude.replicate take t ++ acc+            goNary seen n2 rest acc'+ -- | Creates a singleton trie from an e-class id trie :: EClassId -> IntMap IntTrie -> IntTrie-trie eid = IntTrie (Set.singleton eid)+trie eid = IntTrie {-# INLINE trie #-}  -- | Check whether an e-class is a constant value-isConst :: Monad m => EClassId -> EGraphST m Bool-isConst eid = do ec <- gets ((IntMap.! eid) . _eClass)+isConst :: ClassStore m => EClassId -> EGraphST m Bool+isConst eid = do ec <- getEClass eid                  case (_consts . _info) ec of                    ConstVal _ -> pure True                    _          -> pure False {-# INLINE isConst #-} -getFitness :: Monad m => EClassId -> EGraphST m (Maybe Double)-getFitness c = gets (_fitness . _info . (IntMap.! c) . _eClass)+getFitness :: ClassStore m => EClassId -> EGraphST m (Maybe Double)+getFitness c = (_fitness . _info) <$> getEClass c {-# INLINE getFitness #-}-getTheta :: Monad m => EClassId -> EGraphST m ([PVector])-getTheta c = gets (_theta . _info . (IntMap.! c) . _eClass)+getTheta :: ClassStore m => EClassId -> EGraphST m ([Target])+getTheta c = (_theta . _info) <$> getEClass c {-# INLINE getTheta #-}-getSize :: Monad m => EClassId -> EGraphST m Int-getSize c = gets (_size . _info . (IntMap.! c) . _eClass)+getSize :: ClassStore m => EClassId -> EGraphST m Int+getSize c = (_size . _info) <$> getEClass c {-# INLINE getSize #-} isSizeOf :: (Int -> Bool) -> EClass -> Bool isSizeOf p = p . _size . _info {-# INLINE isSizeOf #-}-getBestFitness :: Monad m => EGraphST m (Maybe Double)+getBestFitness :: ClassStore m => EGraphST m (Maybe Double) getBestFitness = do-    bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)-    gets (_fitness . _info . (IntMap.! bec) . _eClass)-getDL :: Monad m => EClassId -> EGraphST m (Maybe Double)-getDL c = gets (_dl . _info . (IntMap.! c) . _eClass)+    mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB)+    case mbec of+      Just bec -> (_fitness . _info) <$> getEClass bec+      Nothing  -> pure Nothing+getDL :: ClassStore m => EClassId -> EGraphST m (Maybe Double)+getDL c = (_dl . _info) <$> getEClass c {-# INLINE getDL #-}

@@ -15,28 +15,24 @@ module Algorithm.EqSat.Info where  import Control.Lens ( over )-import Control.Monad --(forM, forM_, when, foldM, void)+import Control.Monad import Control.Monad.State import Data.AEq (AEq ((~==)))-import Data.IntMap (IntMap) -- , delete, empty, insert, toList)+import Data.IntMap (IntMap) import qualified Data.IntMap as IntMap import Data.Map (Map) import qualified Data.Map as Map import Data.SRTree-import Data.SRTree.Eval (evalFun, evalOp, PVector)+import Data.SRTree.Eval (evalFun, evalOp, Target) import Data.HashSet (HashSet) import qualified Data.HashSet as Set+import qualified Data.Set as RangeSet import qualified Data.IntSet as IntSet import Algorithm.EqSat.Egraph-import Data.AEq (AEq ((~==))) import Algorithm.EqSat.Queries -import Data.Maybe import qualified Data.Set as TrueSet-import Data.Sequence (Seq(..), (><)) -import Debug.Trace- -- * Data related functions   -- | join data from two e-classes@@ -84,51 +80,71 @@     combineConsts (ConstVal x) (ParamIx ix) = ConstVal x -- p - p = 0     combineConsts x y = error (show x <> " " <> show y) +-- | Fetch consts, cost, and size for all children in a single state traversal+getChildrenData :: ClassStore m => [EClassId] -> EGraphST m [(Consts, Cost, Int)]+getChildrenData ids = do+  ids' <- mapM canonical ids+  mapM (\cid -> do+            ec <- getEClass cid+            let d = _info ec+            pure (_consts d, _cost d, _size d)) ids'+{-# INLINE getChildrenData #-}+ -- | Calculate e-node data (constant values and cost)-makeAnalysis :: Monad m => CostFun -> ENode -> EGraphST m EClassData+makeAnalysis :: ClassStore m => CostFun -> ENode -> EGraphST m EClassData makeAnalysis costFun enode =-  do consts <- calculateConsts enode+  do let cs = eChildren enode+     childData <- getChildrenData cs+     let (consts', costs', sizes) = unzip3 childData+         consts = combineNode enode consts'+         cost   = costNode enode costs'+         sz     = sum sizes      enode' <- canonize enode-     cost   <- calculateCost costFun enode'-     sz <- sum <$> mapM (\ecId -> gets (_size . _info . (IntMap.! ecId) . _eClass)) (childrenOf enode')-     pure $ EData cost enode' consts Nothing Nothing [] (sz+1)+     pure $ EData cost enode' consts Nothing Nothing [] (sz + 1)+  where+    -- ENAry folds children pairwise (constant folding over a multiset); the+    -- binary skeleton cannot represent n children.+    combineNode (ENAry op _) cs = foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) cs+    combineNode _             cs = combineConsts (replaceChildren cs (fromENode enode))+    -- ENAry is a single flattened op node: op cost + sum of child costs.+    costNode (ENAry op _) cs = costFun (Bin (toOp op) 0 0) + sum cs+    costNode _             cs = costFun (replaceChildren cs (fromENode enode)) -getChildrenMinHeight :: Monad m => ENode -> EGraphST m Int+getChildrenMinHeight :: ClassStore m => ENode -> EGraphST m Int getChildrenMinHeight enode = do-  let children = childrenOf enode-      minimum' [] = 0-      minimum' xs = minimum xs-  minimum' <$> mapM (\ec -> gets (_height . (IntMap.! ec) . _eClass)) children+  let children = eChildren enode+  if null children then pure 0 else do+    children' <- mapM canonical children+    hs <- mapM (fmap _height . getEClass) children'+    pure (minimum hs)  -- | update the heights of each e-class -- won't work if there's no root-calculateHeights :: Monad m => EGraphST m ()+calculateHeights :: ClassStore m => EGraphST m () calculateHeights =   do queue   <- findRootClasses-     classes <- gets (Prelude.map fst . IntMap.toList . _eClass)+     classes <- allKeys      let nClasses = length classes      forM_ classes (setHeight nClasses) -- set all heights to max possible height (number of e-classes)      forM_ queue (setHeight 0)          -- set root e-classes height to zero      go queue (TrueSet.fromList queue) 1    -- next height is 1   where+    setHeight :: ClassStore m => Int -> EClassId -> EGraphST m ()     setHeight x eId' =       do eId <- canonical eId'          ec <- getEClass eId          let ec' = over height (const x) ec-         modify' $ over eClass (IntMap.insert eId ec')+         insertClass ec' +    setMinHeight :: ClassStore m => Int -> EClassId -> EGraphST m ()     setMinHeight x eId' = -- set height to the minimum between current and x       do eId <- canonical eId'          h <- _height <$> getEClass eId          setHeight (min h x) eId -    getChildrenEC :: Monad m => EClassId -> EGraphST m [EClassId]-    getChildrenEC ec' = do ec <- canonical ec'-                           gets (concatMap childrenOf' . _eNodes . (IntMap.! ec) . _eClass)--    childrenOf' (_, -1, -1, _) = []-    childrenOf' (_, e1, -1, _) = [e1]-    childrenOf' (_, e1, e2, _) = [e1, e2]+    getChildrenEC :: ClassStore m => EClassId -> EGraphST m [EClassId]+    getChildrenEC ec' = do ec <- getEClass ec'+                           pure $ concatMap eChildren (_eNodes ec)      go [] _    _ = pure ()     go qs tabu h =@@ -138,19 +154,23 @@          go childrenL (TrueSet.union tabu childrenOf) (h+1) -- move one breadth search style  -- | calculates the cost of a node-calculateCost :: Monad m => CostFun -> SRTree EClassId -> EGraphST m Cost-calculateCost f t =-  do let cs = childrenOf t+calculateCost :: ClassStore m => CostFun -> ENode -> EGraphST m Cost+calculateCost f enode =+  do let cs = eChildren enode      costs <- traverse (fmap (_cost . _info) . getEClass) cs-     pure . f $ replaceChildren costs t+     pure $ case enode of+              ENAry op _ -> f (Bin (toOp op) 0 0) + sum costs+              _          -> f (replaceChildren costs (fromENode enode))  -- | check whether an e-node evaluates to a const-calculateConsts :: Monad m => SRTree EClassId -> EGraphST m Consts-calculateConsts t =-  do let cs = childrenOf t-     eg <- get+calculateConsts :: ClassStore m => ENode -> EGraphST m Consts+calculateConsts enode =+  do let cs = eChildren enode      consts <- traverse (fmap (_consts . _info) . getEClass) cs-     case combineConsts $ replaceChildren consts t of+     let c = case enode of+               ENAry op _ -> foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) consts+               _          -> combineConsts (replaceChildren consts (fromENode enode))+     case c of           ConstVal x | isNaN x -> pure (ConstVal x)           a -> pure a @@ -168,41 +188,35 @@     evalOp' (ConstVal x) (ConstVal y) = ConstVal $ evalOp op x y     evalOp' _            _            = NotConst -insertFitness :: Monad m => EClassId -> Double -> [PVector] -> EGraphST m ()-insertFitness eId' fit params = do-  eId <- canonical eId'-  tree <- getBestExpr' eId-  let p = fromIntegral (length params)-  let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0-  ec <- gets ((IntMap.! eId) . _eClass)-  let oldFit  = _fitness . _info $ ec-  --when (oldFit < Just fit) $ do-  let newInfo = (_info ec){_fitness = Just fit, _theta = params}-      newEc   = ec{_info = newInfo}-      sz = _size newInfo-  modify' $ over eClass (IntMap.insert eId newEc)-  if (isNothing oldFit)-    then modify' $ over (eDB . unevaluated) (IntSet.delete eId)-                 . over (eDB . fitRangeDB) (insertRange eId fit)-                 . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)-                 . over (eDB . dlRangeDB) (insertRange eId f_compl)-    else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit))+insertFitness :: ClassStore m => EClassId -> Double -> [Target] -> EGraphST m ()+insertFitness eId' fit params =+  do eId <- canonical eId'+     tree <- getBestExpr eId+     let p = fromIntegral (length params)+     let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0+     ec <- getEClass eId+     let oldFit  = _fitness . _info $ ec+     let newInfo = (_info ec){_fitness = Just fit, _theta = params}+         newEc   = ec{_info = newInfo}+         sz = _size newInfo+     insertClass newEc+     case oldFit of+       Nothing -> modify' $ over (eDB . unevaluated) (IntSet.delete eId)+                    . over (eDB . fitRangeDB) (insertRange eId fit)+                    . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty)+                    . over (eDB . dlRangeDB) (insertRange eId f_compl)+       Just oldVal -> modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId oldVal)+                                 . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit . removeRange eId oldVal) sz) -insertDL :: Monad m => EClassId -> Double -> EGraphST m ()-insertDL eId fit' = do-  let fit = negate fit'-  ec <- gets ((IntMap.! eId) . _eClass)-  let sz = _size . _info $ ec-      newInfo = (_info ec){_dl = Just fit'}-      newEc   = ec{_info=newInfo}-  modify' $ over eClass (IntMap.insert eId newEc)-  modify' $ over (eDB . dlRangeDB) (insertRange eId fit)-          . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)+insertDL :: ClassStore m => EClassId -> Double -> EGraphST m ()+insertDL eId fit' =+  do let fit = negate fit'+     ec <- getEClass eId+     let sz = _size . _info $ ec+         newInfo = (_info ec){_dl = Just fit'}+         newEc   = ec{_info=newInfo}+     insertClass newEc+     modify' $ over (eDB . dlRangeDB) (insertRange eId fit)+             . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty) --- | TODO: remove from here gets the best expression given the default cost function-getBestExpr' :: Monad m => EClassId -> EGraphST m (Fix SRTree)-getBestExpr' eid = do eid' <- canonical eid-                      best <- gets (_best . _info . (IntMap.! eid') . _eClass)-                      childs <- mapM getBestExpr' $ childrenOf best-                      pure . Fix $ replaceChildren childs best 

@@ -21,66 +21,54 @@ import qualified Data.IntMap as IntMap import qualified Data.Map as Map import qualified Data.HashSet as Set+import qualified Data.Set as RangeSet import Control.Monad.State ( gets, modify' )-import Control.Monad ( filterM ) import Control.Lens ( over ) import Data.Maybe-import Data.Sequence ( Seq(..) )-import qualified Data.Sequence as FingerTree-import qualified Data.Foldable as Foldable import Data.SRTree (childrenOf) -import Debug.Trace---- this is too slow for now, it needs a db of its own--- basically a db for each query we need-getEClassesThat :: Monad m => (EClass -> Bool) -> EGraphST m [EClassId]+getEClassesThat :: ClassStore m => (EClass -> Bool) -> EGraphST m [EClassId] getEClassesThat p = do-    gets (map fst . filter (\(ecId, ec) -> p ec) . IntMap.toList . _eClass)-    --go ecs-        where-            go :: Monad m => [EClassId] -> EGraphST m [EClassId]-            go [] = pure []-            go (ecId:ecs) = do ec <- gets (p . (IntMap.! ecId) . _eClass)-                               ecs' <- go ecs-                               if ec-                                  then pure (ecId:ecs')-                                  else pure ecs'+    classes <- allClasses+    pure [ _eClassId ec | ec <- classes, p ec ] -updateFitness :: Monad m => Double -> EClassId -> EGraphST m ()+updateFitness :: ClassStore m => Double -> EClassId -> EGraphST m () updateFitness f ecId = do-   ec   <- gets ((IntMap.! ecId) . _eClass)+   ec   <- getEClass ecId    let info = _info ec-   modify' $ over eClass (IntMap.insert ecId ec{_info=info{_fitness = Just f}})+   insertClass ec{_info=info{_fitness = Just f}}  -- | returns all the root e-classes (e-class without parents)-findRootClasses :: Monad m => EGraphST m [EClassId]-findRootClasses = gets (Prelude.map fst . Prelude.filter isParent . IntMap.toList . _eClass)+findRootClasses :: ClassStore m => EGraphST m [EClassId]+findRootClasses = do+    classes <- allClasses+    pure [ _eClassId ec | ec <- classes, isParent (_eClassId ec, ec) ]   where     isParent (k, v) = Prelude.null (_parents v) ||  (k `Set.member` (Set.map fst (_parents v)))  -- | returns the e-class id with the best fitness that -- is true to a predicate-getTopECLassThat :: Monad m => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopECLassThat :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopECLassThat b n p = do   let f = if b then _fitRangeDB else _dlRangeDB   gets (f . _eDB)     >>= go n []   where-    go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+    go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]     go 0 bests rt = pure bests-    go m bests rt = case rt of-                       Empty   -> pure bests-                       t :|> y -> do let x = snd y-                                     ecId <- canonical x-                                     ec <- gets ((IntMap.! ecId) . _eClass)-                                     if (isInfinite . fromJust . _fitness . _info $ ec)-                                       then go m bests t-                                       else if p ec-                                              then go (m-1) (ecId:bests) t-                                              else go m bests t+    go m bests rt = case RangeSet.maxView rt of+                       Nothing -> pure bests+                       Just (y, t) ->+                         let x = snd y+                         in do ecId <- canonical x+                               ec <- getEClass ecId+                               if (maybe True (isInfinite) . _fitness . _info $ ec)+                                 then go m bests t+                                 else if p ec+                                   then go (m-1) (ecId:bests) t+                                   else go m bests t -getTopEClassInRange :: Monad m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId]+getTopEClassInRange :: ClassStore m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId] getTopEClassInRange b n p range = do   let f = if b then _fitRangeDB else _dlRangeDB   gets (f . _eDB)@@ -92,43 +80,45 @@       | v > y = 1       | otherwise = 1  -    go :: Monad m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId]+    go :: ClassStore m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId]     go _ bests []      _ = pure bests      go 0 bests (r:rs) rt = go n bests rs rt-    go m bests (r:rs) rt = case rt of-                             Empty   -> pure bests-                             t :|> y -> do let x = snd y-                                           ecId <- canonical x-                                           ec <- gets ((IntMap.! ecId) . _eClass)-                                           if (isInfinite . fromJust . _fitness . _info $ ec)-                                             then go m bests (r:rs) t-                                             else do let v = p ec -                                                     case (v `inRange` r) of-                                                       0  -> go (m-1) (ecId:bests) (r:rs) t -- it is in range, go to the next range -                                                       -1 -> go n bests rs (t :|> y) -- it is smaller than the range, get the first n of the next range-                                                       1  -> go m bests (r:rs) t -- y is still greater than the range, keep looking in the same range+    go m bests (r:rs) rt = case RangeSet.maxView rt of+                             Nothing -> pure bests+                             Just (y, t) ->+                               let x = snd y+                               in do ecId <- canonical x+                                     ec <- getEClass ecId+                                     if (maybe True (isInfinite) . _fitness . _info $ ec)+                                       then go m bests (r:rs) t+                                       else do let v = p ec+                                               case (v `inRange` r) of+                                                 0  -> go (m-1) (ecId:bests) (r:rs) t+                                                 -1 -> go n bests rs (RangeSet.insert y t)+                                                 1  -> go m bests (r:rs) t -getTopECLassIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopECLassIn :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopECLassIn b n p ecs' = do   let f = if b then _fitRangeDB else _dlRangeDB   gets (f . _eDB)     >>= go n []   where     ecs = Set.fromList ecs'-    go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+    go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]     go 0 bests rt = pure bests-    go m bests rt = case rt of-                       Empty   -> pure bests-                       t :|> y -> do let x = snd y-                                     ecId <- canonical x-                                     ec <- gets ((IntMap.! ecId) . _eClass)-                                     if (isInfinite . fromJust . _fitness . _info $ ec)-                                       then go m bests t -- pure bests-                                       else if ecId `Set.member` ecs && p ec-                                              then go (m-1) (ecId:bests) t-                                              else go m bests t+    go m bests rt = case RangeSet.maxView rt of+                       Nothing -> pure bests+                       Just (y, t) ->+                         let x = snd y+                         in do ecId <- canonical x+                               ec <- getEClass ecId+                               if (maybe True (isInfinite) . _fitness . _info $ ec)+                                 then go m bests t+                                 else if ecId `Set.member` ecs && p ec+                                   then go (m-1) (ecId:bests) t+                                   else go m bests t -getTopECLassNotIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopECLassNotIn :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopECLassNotIn b n p ecs' = do   let f = if b then _fitRangeDB else _dlRangeDB   gets (f . _eDB)@@ -136,65 +126,65 @@   where     ecs = Set.fromList ecs' -    go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+    go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]     go 0 bests rt = pure bests-    go m bests rt = case rt of-                       Empty   -> pure bests-                       t :|> y -> do let x = snd y-                                     ecId <- canonical x-                                     ec <- gets ((IntMap.! ecId) . _eClass)-                                     if (isInfinite . fromJust . _fitness . _info $ ec)-                                       then go m bests t-                                       else if not (ecId `Set.member` ecs) && p ec-                                              then go (m-1) (ecId:bests) t-                                              else go m bests t+    go m bests rt = case RangeSet.maxView rt of+                       Nothing -> pure bests+                       Just (y, t) ->+                         let x = snd y+                         in do ecId <- canonical x+                               ec <- getEClass ecId+                               if (maybe True (isInfinite) . _fitness . _info $ ec)+                                 then go m bests t+                                 else if not (ecId `Set.member` ecs) && p ec+                                   then go (m-1) (ecId:bests) t+                                   else go m bests t -getAllEvaluatedEClasses :: Monad m => EGraphST m [EClassId]+getAllEvaluatedEClasses :: ClassStore m => EGraphST m [EClassId] getAllEvaluatedEClasses = do   gets (_fitRangeDB . _eDB)     >>= go []   where-    go :: Monad m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId]-    go bests rt = case rt of-                    Empty   -> pure bests-                    t :|> y -> do let x = snd y-                                  ecId <- canonical x-                                  ec <- gets ((IntMap.! ecId) . _eClass)-                                  if (isInfinite . fromJust . _fitness . _info $ ec)-                                    then go bests t-                                    else go (ecId:bests) t+    go :: ClassStore m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+    go bests rt = case RangeSet.maxView rt of+                    Nothing -> pure bests+                    Just (y, t) ->+                      let x = snd y+                      in do ecId <- canonical x+                            ec <- getEClass ecId+                            if (maybe True (isInfinite) . _fitness . _info $ ec)+                              then go bests t+                              else go (ecId:bests) t  getTopEClassWithSize :: Monad m => Bool -> Int -> Int -> EGraphST m [EClassId] getTopEClassWithSize b sz n = do    let fun = if b then _sizeFitDB else _sizeDLDB    gets (go n [] . (IntMap.!? sz) . fun . _eDB)-    -- >>= mapM canonical   where-    -- go :: Monad m => Int -> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId]     go _ bests Nothing   = []     go 0 bests (Just rt) = bests-    go m bests (Just rt) = case rt of-                             Empty   -> bests-                             t :|> (f, x) -> if isInfinite f || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t)+    go m bests (Just rt) = case RangeSet.maxView rt of+                             Nothing         -> bests+                             Just ((f, x), t) -> if isInfinite f || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t) -getTopFitEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopFitEClassThat :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopFitEClassThat  = getTopECLassThat True-getTopDLEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopDLEClassThat :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopDLEClassThat   = getTopECLassThat False-getTopFitEClassIn :: Monad m =>  Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassIn :: ClassStore m =>  Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopFitEClassIn    = getTopECLassIn True-getTopDLEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopDLEClassIn     = getTopECLassIn False-getTopFitEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopFitEClassNotIn = getTopECLassNotIn True-getTopDLEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]-getTopDLEClassNotIn  = getTopECLassNotIn True+getTopDLEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassNotIn  = getTopECLassNotIn False getTopFitEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId] getTopFitEClassWithSize = getTopEClassWithSize True getTopDLEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId] getTopDLEClassWithSize  = getTopEClassWithSize False -rebuildAllRanges :: Monad m => EGraphST m ()+rebuildAllRanges :: ClassStore m => EGraphST m () rebuildAllRanges = do szF <- gets (_sizeFitDB._eDB) >>= traverse rebuildRange                       dlF <- gets (_sizeDLDB._eDB) >>= traverse rebuildRange                       fR  <- gets (_fitRangeDB._eDB) >>= rebuildRange@@ -205,17 +195,19 @@                               . over (eDB.sizeFitDB) (const szF)                               . over (eDB.sizeDLDB) (const dlF) -canonizeRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)-canonizeRange = traverse (\(x, eid) -> (x,) <$> canonical eid)+canonizeRange :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double)+canonizeRange = fmap RangeSet.fromList . mapM (\(x, eid) -> (x,) <$> canonical eid) . RangeSet.toList -rebuildRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)-rebuildRange rt = go Set.empty Empty <$> canonizeRange rt+rebuildRange :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double)+rebuildRange rt = do+  canonRt <- canonizeRange rt+  pure $ snd $ go canonRt   where-    go :: Set.HashSet EClassId -> RangeTree Double -> RangeTree Double -> RangeTree Double-    go seen root Empty = root-    go seen root (xs :|> (x,eid)) = go (Set.insert eid seen)-                                       (if Set.member eid seen-                                          then root-                                          else (x, eid) :<| root)-                                        xs -- (Prelude.filter ((/= eid) . snd) xs)+    go rt' = case RangeSet.maxView rt' of+               Nothing -> (Set.empty, RangeSet.empty)+               Just ((x, eid), rest) ->+                 let (seen, result) = go rest+                 in if Set.member eid seen+                      then (seen, result)+                      else (Set.insert eid seen, RangeSet.insert (x, eid) result) 

@@ -15,26 +15,33 @@  import Data.SRTree import Data.SRTree.Datasets+import Data.SRTree.Eval (compileLoss) import System.Random import Control.Monad.State.Strict+import Control.Concurrent (getNumCapabilities)+import Control.Concurrent.Async (mapConcurrently)+import Data.Maybe (catMaybes)+import Control.Exception (evaluate)+import qualified Control.DeepSeq as DeepSeq import Algorithm.EqSat.Egraph import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.AD (ADBackEnd(..))+import Algorithm.SRTree.AD.Unboxed (setMTPopParallel) import qualified Data.IntMap as IM import qualified Data.IntSet as IntSet import qualified Data.SRTree.Random as Random import Data.Function ( on )-import Algorithm.SRTree.Likelihoods import Algorithm.SRTree.NonlinearOpt import Control.Monad ( when, replicateM, forM, forM_ )-import Algorithm.EqSat.Egraph-import Algorithm.SRTree.Opt+import Numeric.Optimization.NLOPT import Algorithm.EqSat.Info import Algorithm.EqSat.Build-import Data.Maybe ( fromJust ) import Data.SRTree.Random import Algorithm.EqSat.Queries import Data.List ( maximumBy )-import qualified Data.Map.Strict as Map+import qualified Data.List as Data.List+import qualified Data.HashMap.Strict as HashMap+import qualified Data.Vector.Unboxed as V  -- Environment of an e-graph with support to random generator and IO type RndEGraph a = EGraphST (StateT StdGen IO) a@@ -46,6 +53,55 @@ rnd = lift {-# INLINE rnd #-} +-- | Run an 'RndEGraph' action against a read-only egraph snapshot with the given+-- generator (for concurrent workers that do not mutate the shared egraph).+runRndEGraph :: EGraph -> StdGen -> RndEGraph a -> IO a+runRndEGraph eg g m = do+  ((a, _), _) <- runStateT (runStateT m eg) g+  pure a+{-# INLINE runRndEGraph #-}++-- | Fit a batch of e-classes in parallel, then insert the results serially.+-- Semantics mirror 'updateIfNothing' (skip already-fitted) unless 'force' is+-- True. The shared 'StdGen' is split once; each worker gets its own generator,+-- so the global draw sequence differs from the serial search (acceptable).+-- While the batch runs, the MultiThread backend is switched to single-chunk so+-- cores go to the batch rather than oversubscribing the inner per-tree split.+fitBatch :: Bool+         -> (Fix SRTree -> RndEGraph (Double, [Target]))+         -> [EClassId]+         -> RndEGraph ()+fitBatch force fitFun ecs0 = do+  ecs <- Prelude.mapM canonical ecs0+  jobs <- fmap catMaybes $ forM ecs $ \ec -> do+            mf <- getFitness ec+            if force || mf == Nothing+               then do tree <- getBestExpr ec+                       pure (Just (ec, tree))+               else pure Nothing+  case jobs of+    [] -> pure ()+    _  -> do+      nCaps <- io getNumCapabilities+      g0 <- rnd get+      let (seed, g1) = random g0 :: (Int, StdGen)+          gs    = [ mkStdGen (seed + fromIntegral i) | i <- [0 .. length jobs - 1] ]+          jobsG = [ (ec, tree, g) | ((ec, tree), g) <- zip jobs gs ]+          chunk k xs = [ [ xs !! j | j <- [i, i + k .. length xs - 1] ] | i <- [0 .. k - 1] ]+      rnd (put g1)+      eg <- get+      io (setMTPopParallel True)+      results <- io $ fmap concat (mapConcurrently (mapM (runJob eg fitFun)) (chunk nCaps jobsG))+      io (setMTPopParallel False)+      forM_ results $ \(ec0, f, p) -> insertFitness ec0 f p+  where+    runJob :: EGraph -> (Fix SRTree -> RndEGraph (Double, [Target])) -> (EClassId, Fix SRTree, StdGen) -> IO (EClassId, Double, [Target])+    runJob eg fit' (ec, tree, g) = do+      (f, p) <- runRndEGraph eg g (fit' tree)+      f' <- evaluate (DeepSeq.force f)+      p' <- evaluate (DeepSeq.force p)+      pure (ec, f', p')+ myCost :: SRTree Int -> Int myCost (Var _)     = 1 myCost (Const _)   = 1@@ -59,34 +115,34 @@                               while p arg' prog                       else pure arg -fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)-fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =-  if isNaN val -- || isNaN tr-    then (-(1/0), theta) -- infinity+fitnessFun :: ADBackEnd -> Bool -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> Target -> (Double, Target)+fitnessFun backend skipVal nIter loss (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =+  if isNaN val+    then (-(1/0), theta)     else (val, theta)   where-    --tree          = relabelParams _tree-    nParams       = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig-    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta-    --tr            = evalF x y mYErr-    val           = evalF x_val y_val mYErr_val+    nParams       = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL Gaussian then 1 else 0+    (theta, lossVal, _) = minimizeNLL' VAR1 backend loss mYErr nIter x y tree thetaOrig+    evalF a b c   = negate $ compileLoss a (buildLoss loss (fromIntegral (V.length b)) tree) b c $ if nParams == 0 then thetaOrig else theta+    -- at folds=1 the validation split is the training data itself, so the+    -- train loss returned by minimizeNLL' already is the val loss; skipping+    -- the separate compileLoss below avoids re-evaluating every expression.+    val           = if skipVal then negate lossVal else evalF x_val y_val mYErr_val  --{-# INLINE fitnessFun #-} -fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)-fitnessFunRep nRep nIter distribution dataTrain dataVal tree = do-    let nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+fitnessFunRep :: ADBackEnd -> Bool -> Int -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, Target)+fitnessFunRep backend skipVal nRep nIter loss dataTrain dataVal tree = do+    let nParams = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL Gaussian then 1 else 0     thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)-    let fits = maximumBy (compare `on` fst) $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal tree) thetaOrigs-    pure fits+    pure $ maximumBy (\(x, _) (y, _) -> compare x y) $ Prelude.map (fitnessFun backend skipVal nIter loss dataTrain dataVal tree) thetaOrigs --{-# INLINE fitnessFunRep #-}  -fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [PVector])-fitnessMV shouldReparam nRep nIter distribution dataTrainsVals _tree = do+fitnessMV :: ADBackEnd -> Bool -> Bool -> Int -> Int -> Loss -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [Target])+fitnessMV backend skipVal shouldReparam nRep nIter loss dataTrainsVals _tree = do   let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree-  response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv tree+  response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep backend skipVal nRep nIter loss dt dv tree   pure (minimum (Prelude.map fst response), Prelude.map snd response)  @@ -95,7 +151,7 @@  -- RndEGraph utils -- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val-insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [Target])) -> RndEGraph EClassId insertExpr t fitFun = do     ecId <- fromTree myCost t >>= canonical     (f, p) <- fitFun t@@ -116,29 +172,28 @@ pickRndSubTree :: RndEGraph (Maybe EClassId) pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)                     if not (null ecIds)-                          then do rndId' <- rnd $ randomFrom ecIds-                                  rndId  <- canonical rndId'-                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)-                                  case constType of-                                    NotConst -> pure $ Just rndId-                                    _        -> pure Nothing-                          else pure Nothing+                      then do rndId' <- rnd $ randomFrom ecIds+                              rndId  <- canonical rndId'+                              constType <- (_consts . _info) <$> getEClass rndId+                              case constType of+                                NotConst -> pure $ Just rndId+                                _        -> pure Nothing+                      else pure Nothing  getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n) getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n)  getBestExprWithSize n =         do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical-           if (not (null ec))-            then do-              bestFit <- getFitness $ head ec-              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)-              pure [(head ec, bestFit)]-            else pure []+           case ec of+             (x:_) -> do bestFit <- getFitness x+                         bestP   <- (_theta . _info) <$> getEClass x+                         pure [(x, bestFit)]+             []    -> pure []  insertRndExpr maxSize rndTerm rndNonTerm =       do grow <- rnd toss-         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. max 1 maxSize])          t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow          fromTree myCost t >>= canonical @@ -152,11 +207,11 @@  --printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph () printBest fitFun printExprFun = do-      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical-      bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)-      --refit fitFun bec-      --io.print $ "should be " <> show bestFit-      printExprFun 0 bec+      mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB)+      case mbec of+        Just bec -> do bestFit <- (_fitness . _info) <$> getEClass bec+                       printExprFun 0 bec+        Nothing  -> pure ()  --paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph () paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))@@ -166,18 +221,17 @@         | n > maxSize = pure []         | otherwise   = do             ecList <- getBestExprWithSize n-            if not (null ecList)-                then do let (ec, mf) = head ecList-                            f' = fromJust mf-                            improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f'-                        ec' <- canonical ec-                        if improved-                                then do refit fitFun ec'-                                        t <- printExprFun ix ec'-                                        ts <- go (n+1) (ix + if improved then 1 else 0) (max f f')-                                        pure (t:ts)-                                else go (n+1) (ix + if improved then 1 else 0) (max f f')-                else go (n+1) ix f+            case ecList of+              ((ec, Just f'):_) -> do+                let improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f'+                ec' <- canonical ec+                if improved+                  then do refit fitFun ec'+                          t <- printExprFun ix ec'+                          ts <- go (n+1) (ix + if improved then 1 else 0) (max f f')+                          pure (t:ts)+                  else go (n+1) (ix + if improved then 1 else 0) (max f f')+              _ -> go (n+1) ix f  evaluateUnevaluated fitFun = do           ec <- gets (IntSet.toList . _unevaluated . _eDB)@@ -196,26 +250,26 @@  -- | check whether an e-node exists or does not exist in the e-graph doesExist, doesNotExist :: ENode -> RndEGraph Bool-doesExist en = gets ((Map.member en) . _eNodeToEClass)-doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)+doesExist en = gets ((HashMap.member en) . _eNodeToEClass)+doesNotExist en = gets ((not . HashMap.member en) . _eNodeToEClass)  -- | check whether the partial tree defined by a list of ancestors will create -- a non-existent expression when combined with a certain e-node. doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool-doesNotExistGens []              en = gets ((Map.notMember en) . _eNodeToEClass)-doesNotExistGens (mGrand:grands) en = do  b <- gets ((Map.notMember en) . _eNodeToEClass)+doesNotExistGens []              en = gets ((not . HashMap.member en) . _eNodeToEClass)+doesNotExistGens (mGrand:grands) en = do  b <- gets ((not . HashMap.member en) . _eNodeToEClass)                                           if b                                             then pure True                                             else case mGrand of                                                 Nothing -> pure False-                                                Just gf -> do ec  <- gets ((Map.! en) . _eNodeToEClass)+                                                Just gf -> do ec  <- gets ((HashMap.! en) . _eNodeToEClass)                                                               en' <- canonize (gf ec)                                                               doesNotExistGens grands en'  -- | check whether combining a partial tree `parent` with the e-node `en'` -- will create a new expression checkToken parent en' = do  en <- canonize en'-                            mEc <- gets ((Map.!? en) . _eNodeToEClass)+                            mEc <- gets (HashMap.lookup en . _eNodeToEClass)                             case mEc of                                 Nothing -> pure True                                 Just ec -> do ec' <- canonical ec

@@ -1,244 +0,0 @@--------------------------------------------------------------------------------- |--- Module      :  Algorithm.EqSat.Search--- Copyright   :  (c) Fabricio Olivetti 2021 - 2024--- License     :  BSD3--- Maintainer  :  fabricio.olivetti@gmail.com--- Stability   :  experimental--- Portability :------ Support functions for search symbolic expressions with e-graphs-----------------------------------------------------------------------------------module Algorithm.EqSat.SearchSRCache where--import Data.SRTree-import Data.SRTree.Datasets-import System.Random-import Control.Monad.State.Strict-import Algorithm.EqSat.Egraph-import Algorithm.SRTree.Likelihoods-import qualified Data.IntMap as IM-import qualified Data.IntSet as IntSet-import qualified Data.SRTree.Random as Random-import Data.Function ( on )-import Algorithm.SRTree.Likelihoods-import Algorithm.SRTree.NonlinearOpt-import Control.Monad ( when, replicateM, forM, forM_ )-import Algorithm.EqSat.Egraph-import Algorithm.SRTree.Opt-import Algorithm.EqSat.Info-import Algorithm.EqSat.Build-import Data.Maybe ( fromJust )-import Data.SRTree.Random-import Algorithm.EqSat.Queries-import Data.List ( maximumBy )-import qualified Data.Map.Strict as Map-import Control.Monad.Identity--import Debug.Trace---- Environment of an e-graph with support to random generator and IO-type RndEGraph a = EGraphST (StateT StdGen (StateT [ECache] IO)) a--io :: IO a -> RndEGraph a-io = lift . lift . lift-{-# INLINE io #-}-getCache :: StateT [ECache] IO a -> RndEGraph a-getCache = lift . lift-rnd :: StateT StdGen (StateT [ECache] IO)  a -> RndEGraph a-rnd = lift-{-# INLINE rnd #-}--myCost :: SRTree Int -> Int-myCost (Var _)     = 1-myCost (Const _)   = 1-myCost (Param _)   = 1-myCost (Bin _ l r) = 2 + l + r-myCost (Uni _ t)   = 3 + t--while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f t-while p arg prog = do if (p arg)-                      then do arg' <- prog arg-                              while p arg' prog-                      else pure arg--fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> EGraph -> EClassId -> ECache -> PVector -> (Double, PVector, ECache)-fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) egraph root cache thetaOrig =-  if isNaN val -- || isNaN tr-    then (-(1/0), theta,cache') -- infinity-    else (val, theta, cache')-  where-    tree          = runIdentity $ getBestExpr root `evalStateT` egraph-    nParams       = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    (theta, val, _, cache') = minimizeNLLEGraph VAR1 distribution mYErr nIter x y egraph root cache thetaOrig-    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta-    -- val           = evalF x_val y_val mYErr_val----{-# INLINE fitnessFun #-}--fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> EClassId -> ECache -> RndEGraph (Double, PVector, ECache)-fitnessFunRep nRep nIter distribution dataTrain dataVal root cache = do-    egraph <- get-    let nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-        fst' (a, _, _) = a-    thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)-    let fits = maximumBy (compare `on` fst') $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal egraph root cache) thetaOrigs-    pure fits---{-# INLINE fitnessFunRep #-}---fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector])-fitnessMV shouldReparam nRep nIter distribution dataTrainsVals root = do-  -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree-  -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm'-  caches <- getCache get-  response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache-  getCache $ put (Prelude.map trd response)-  pure (minimum (Prelude.map fst' response), Prelude.map snd' response)-  where fst' (a, _, _) = a-        snd' (_, a, _) = a-        trd  (_, _, a) = a--fitnessMVNoCache :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector])-fitnessMVNoCache shouldReparam nRep nIter distribution dataTrainsVals root = do-  -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree-  -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm'-  caches <- getCache get-  response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache-  pure (minimum (Prelude.map fst' response), Prelude.map snd' response)-  where fst' (a, _, _) = a-        snd' (_, a, _) = a-        trd  (_, _, a) = a------ RndEGraph utils--- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val-insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId-insertExpr t fitFun = do-    ecId <- fromTree myCost t >>= canonical-    (f, p) <- fitFun t-    insertFitness ecId f p-    pure ecId-  where powabs l r  = Fix (Bin PowerAbs l r)--updateIfNothing fitFun ec = do-      mf <- getFitness ec-      case mf of-        Nothing -> do-          --t <- getBestExpr ec-          (f, p) <- fitFun ec-          insertFitness ec f p-          pure True-        Just _ -> pure False--pickRndSubTree :: RndEGraph (Maybe EClassId)-pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)-                    if not (null ecIds)-                          then do rndId' <- rnd $ randomFrom ecIds-                                  rndId  <- canonical rndId'-                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)-                                  case constType of-                                    NotConst -> pure $ Just rndId-                                    _        -> pure Nothing-                          else pure Nothing--getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)-getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n)--getBestExprWithSize n =-        do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical-           if (not (null ec))-            then do-              bestFit <- getFitness $ head ec-              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)-              pure [(head ec, bestFit)]-            else pure []--insertRndExpr maxSize rndTerm rndNonTerm =-      do grow <- rnd toss-         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])-         t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow-         fromTree myCost t >>= canonical--refit fitFun ec = do-  --t <- getBestExpr ec-  (f, p) <- fitFun ec-  mf <- getFitness ec-  case mf of-    Nothing -> insertFitness ec f p-    Just f' -> when (f > f') $ insertFitness ec f p----printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-printBest fitFun printExprFun = do-      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical-      bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)-      --refit fitFun bec-      --io.print $ "should be " <> show bestFit-      printExprFun 0 bec----paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))-    where-    go :: Int -> Int -> Double -> RndEGraph [[String]]-    go n ix f-        | n > maxSize = pure []-        | otherwise   = do-            ecList <- getBestExprWithSize n-            if not (null ecList)-                then do let (ec, mf) = head ecList-                            f' = fromJust mf-                            improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f'-                        ec' <- canonical ec-                        if improved-                                then do refit fitFun ec'-                                        t <- printExprFun ix ec'-                                        ts <- go (n+1) (ix + if improved then 1 else 0) (max f f')-                                        pure (t:ts)-                                else go (n+1) (ix + if improved then 1 else 0) (max f f')-                else go (n+1) ix f--evaluateUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          forM_ ec $ \c -> do-              --t <- getBestExpr c-              (f, p) <- fitFun c-              insertFitness c f p--evaluateRndUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          c <- rnd . randomFrom $ ec-          --t <- getBestExpr c-          (f, p) <- fitFun c-          insertFitness c f p-          pure c---- | check whether an e-node exists or does not exist in the e-graph-doesExist, doesNotExist :: ENode -> RndEGraph Bool-doesExist en = gets ((Map.member en) . _eNodeToEClass)-doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)---- | check whether the partial tree defined by a list of ancestors will create--- a non-existent expression when combined with a certain e-node.-doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool-doesNotExistGens []              en = gets ((Map.notMember en) . _eNodeToEClass)-doesNotExistGens (mGrand:grands) en = do  b <- gets ((Map.notMember en) . _eNodeToEClass)-                                          if b-                                            then pure True-                                            else case mGrand of-                                                Nothing -> pure False-                                                Just gf -> do ec  <- gets ((Map.! en) . _eNodeToEClass)-                                                              en' <- canonize (gf ec)-                                                              doesNotExistGens grands en'---- | check whether combining a partial tree `parent` with the e-node `en'`--- will create a new expression-checkToken parent en' = do  en <- canonize en'-                            mEc <- gets ((Map.!? en) . _eNodeToEClass)-                            case mEc of-                                Nothing -> pure True-                                Just ec -> do ec' <- canonical ec-                                              ec'' <- canonize (parent ec')-                                              not <$> doesExist ec''

@@ -18,24 +18,35 @@ import Algorithm.EqSat.Egraph import Algorithm.EqSat.DB   ( ClassOrVar,-    Pattern (Fixed, VarPat),+    Condition (Condition),+    NChild (Ch, MapP, Rest),+    Pattern (Fixed, Hole, NAry, VarPat),     Rule (..),+    Subst,+    SubVal (SVMap, SVOne),     getInt,   ) import Control.Monad.State.Strict (evalState)-import Data.IntMap (IntMap)-import qualified Data.IntMap as IM+import Data.IntMap.Strict (IntMap)+import qualified Data.IntMap.Strict as IM import Data.Map (Map) import qualified Data.Map as Map import Data.SRTree -type ConstrFun = Pattern -> Map ClassOrVar ClassOrVar -> EGraph -> Bool +-- | A constraint over a match's substitution: when applied to a substitution it+-- runs in the e-graph monad and fetches e-class data through 'ClassStore', so it+-- works on a paged (out-of-core) graph whose resident cache is bounded/empty.+type ConstrFun = Pattern -> Condition -constrainOnVal :: (Consts -> Bool) -> Pattern -> Map ClassOrVar ClassOrVar -> EGraph -> Bool -constrainOnVal f (VarPat c) subst eg =-    let cid = getInt $ subst Map.! Right (fromEnum c)-     in f (_consts . _info $ _eClass eg IM.! cid)-constrainOnVal _ _ _ _ = False +constrainOnVal :: (Consts -> Bool) -> Pattern -> Condition+constrainOnVal f (VarPat c) = Condition $ \subst -> do+    let cid = getInt $ case Map.lookup (Right (fromEnum c)) subst of+                        Nothing -> error $ "CONSTRAINVAL_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst)+                        Just (SVOne v) -> v+                        Just (SVMap _) -> error $ "CONSTRAINVAL_REST_AS_SINGLE var=" <> show (fromEnum c)+    ec <- getEClass cid+    pure (f (_consts . _info $ ec))+constrainOnVal _ _ = Condition $ \_ -> pure False  -- TODO: aux functions to avoid repeated pattern in constraint creation  --@@ -94,156 +105,149 @@        ConstVal x -> not (isNaN x || isInfinite x)        _          -> True --- basic algebraic rules +-- | e-class ids bound to a rest variable+restEidsOf :: Char -> Subst -> [EClassId]+restEidsOf c subst = case Map.lookup (Right (fromEnum c)) subst of+                       Just (SVMap m) -> expandedList m+                       _              -> []++-- | every e-class bound to a rest variable holds a valid value+allValidRest :: Char -> Condition+allValidRest c = Condition $ \subst -> do+    let eids = restEidsOf c subst+        validEid eid = getEClass eid >>= \ec ->+            pure $ case _consts . _info $ ec of+                     ConstVal x -> not (isNaN x || isInfinite x)+                     _          -> True+    and <$> mapM validEid eids++-- basic algebraic rules rewriteBasic :: [Rule] rewriteBasic =     [-      "x" * "y" :=> "y" * "x"-    , "x" + "y" :=> "y" + "x"-    --, ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"-    --, (powabs "x" "y") * (powabs "x" "z") :=> powabs "x" ("y" + "x")-    , ("x" + "y") + "z" :=> "x" + ("y" + "z")-    , ("x" + "y") - "z" :=> "x" + ("y" - "z")-    --, ("x" + "y") - "z" :=> "x" + ("y" - "z") -- TODO: check that I don't need that-    , ("x" * "y") * "z" :=> "x" * ("y" * "z")-    , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")-    , "x" - ("y" + "z") :=> ("x" - "y") - "z" -- TODO: check that I don't this-    , "x" - ("y" - "z") :=> ("x" - "y") + "z" -- TODO-    , ("x" * "y") / "z" :=> ("x" / "z") * "y" :| isNotZero "z" -- TODO: inv(x) <=> x^-1 , x/y <=> x*y^-1-    , "x" * ("y" / "z") :=> ("x" / "z") * "y" :| isNotZero "z" -- ^-    , "x" / ("y" * "z") :=> ("x" / "z") / "y" :| isNotZero "z" -- ^ TODO: 0 ^-1 check-    , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"-    , ("w" * "x") - ("z" * "x") :=> ("w" - "z") * "x" -- TODO: handle sub :| isConstPt "w" :| isConstPt "z"-    , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y") -- TODO handle with power :| isConstPt "w" :| isConstPt "z" :| isNotZero "z"-    -- TODO: a + b*y :=> b * (a/b + y) :| isNotZero b-    , (("x" * "y") + ("z" * "w")) :=> "x" * ("y" + ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"-    -- , "a" * (("x" * "y") + ("z" * "w")) :=> ("a" * "x") * ("y" + ("z" / "x") * "w") :| isConstPt "a" :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"-    , (("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"-    , (("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x" :| isConstPt "z"-    , "x" * "x" :=> "x" ** 2 -    , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 -    , "x" ** 2 + "x" * "y" :=> "x" * ("x" + "y")-    -- , "x" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL -    -- , "x" + "y" * "z" :=> "y" * ("x" * "y" ** (-1) + "z") :| isNotZero "y" -- GABRIEL +      -- B7/B8/C5: factor a common term out of a sum of products, and the+      -- reverse (distribute), which make x*(y+z) and x*y+x*z equivalent.+      NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1'])+                , Ch (NAry EMul [Ch "x", Rest '2'])+                , Rest '3' ]+        :=>+      NAry EAdd [ Ch (NAry EMul [ Ch "x"+                                , Ch (NAry EAdd [Rest '1', Rest '2'])+                                ])+                , Rest '3' ]+    , NAry EAdd [ Ch (NAry EMul [ Ch "x"+                                , Ch (NAry EAdd [Rest '1'])+                                ])+                , Rest '2' ]+        :=>+      NAry EAdd [ MapP (NAry EMul [Ch "x", Ch Hole]) '1'+                , Rest '2' ]+    -- C5: x*y - z*x = x*(y - z)+    , NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1'])+                , Ch (NAry EMul [Ch (Fixed (Const (-1))), Ch "x", Ch "z"])+                , Rest '3' ]+        :=>+      NAry EAdd [ Ch (NAry EMul [ Ch "x"+                                , Ch (NAry EAdd [Rest '1', Ch (negate (VarPat 'z'))])+                                ])+                , Rest '3' ]+    -- B1: group duplicate factors into a power (x*x = x^2)+    , NAry EMul [Ch "x", Ch "x"] :=> "x" ** 2+    -- C9: binomial expansion of a closed 2-ary square+    , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2+    -- C10: x^2 + x*y + ... = x*(x + y) + ...+    , NAry EAdd [ Ch (Fixed (Bin Power (VarPat 'x') (Fixed (Const 2))))+                , Ch (NAry EMul [Ch "x", Rest '1'])+                , Rest '2' ]+        :=>+      NAry EAdd [ Ch (NAry EMul [ Ch "x"+                                , Ch (NAry EAdd [Ch "x", Rest '1'])+                                ])+                , Rest '2' ]     ]  -- rules for nonlinear functions  rewritesFun :: [Rule] rewritesFun =     [-      log (exp "x") :==: exp (log "x")-    , log (exp "x")  :=> "x"-    -- , exp (log "x")  :=> "x" -- :| isPositive "x" ??? exp(log(x)), x, log(exp(0))-    , log ("x" * "y") :=> log "x" + log "y" :| isConstPos "x" :| isConstPos "y"-    -- , log ("x" / "y") :=> log "x" - log "y" :| isConstPos "x" :| isConstPos "y"+      log (exp "x")  :=> "x"+    -- C11: log(x*y*z*...) = log x + log y + ...+    , log (NAry EMul [Rest '1']) :=> NAry EAdd [MapP (Fixed (Uni Log Hole)) '1']     , log ("x" ** "y") :=> "y" * log "x"     , log (powabs "x" "y") :=> "y" * log (abs "x")-    --, sqrt ("x" ** "y") :=> "x" ** ("y" / 2) :| isEven "y"-    -- , sqrt ("y" * "x") :=> sqrt "y" * sqrt "x" ---    --, sqrt ("y" / "x") :=> sqrt "y" / sqrt "x"-    , abs ("x" * "y") :=> abs "x" * abs "y" -- :| isConstPt "x"+    -- C12: abs(x*y*z*...) = abs x * abs y * ...+    , abs (NAry EMul [Rest '1']) :=> NAry EMul [MapP (Fixed (Uni Abs Hole)) '1']     , abs ("x" ** "y") :=> abs "x" ** "y"-    , abs ("x" - "y") :=> abs ("y" - "x")-    --, sqrt ("z" * ("x" - "y")) :=> sqrt (negate "z") * sqrt ("y" - "x")-    --, sqrt ("z" * ("x" + "y")) :=> sqrt "z" * sqrt ("x" + "y")     , recip (recip "x") :=> "x" :| isNotZero "x"-    , ("x" * "y") ** "z" :==: ("x" ** "z") * ("y" ** "z") -- :| bothSameSign "x" "y"-    , ("x" * "y") ** "z" :==: ("x" ** "z") * ("y" ** "z") -- :| isInteger "z"-    --, recip "x" :==: "x" ** (-1) -- GABRIEL -    --, "x" / "y" :==: "x" * "y" ** (-1) -- GABRIEL +    -- C13: (x*y*z*...)^w = x^w * y^w * ...   [was disabled: combinatorial blowup on (x*x)^t; the multiset matcher + matchCap bound that]+    , (NAry EMul [Rest '1']) ** "z" :=> NAry EMul [MapP (Hole ** VarPat 'z') '1']     , abs "x" ** "y" :=> "x" ** "y" :| isEven "y"-    , sqrt ("x" * "x") :=> abs "x"+    -- C14: sqrt(x*x) = abs x+    , sqrt (NAry EMul [Ch "x", Ch "x"]) :=> abs "x"     ]  -- Rules that reduces redundant parameters constReduction :: [Rule] constReduction =     [-      0 + "x" :=> "x"-    -- , "x" - 0 :=> "x"-    --, 1 * "x" :=> "x"-    -- , 0 / "x" :=> 0 :| isNotZero "x"-    --, "x" - "x" :=> 0 :| isNotParam "x"-    --, "x" / "x" :=> 1 :| isNotZero "x" :| isNotParam "x"+      -- B3: 0 + rest = rest+      NAry EAdd [Ch (Fixed (Const 0)), Rest '1'] :=> NAry EAdd [Rest '1']     , "x" ** 1 :=> "x"     , powabs "x" 1 :=> abs "x" -    -- , "x" * (1 / "x") :=> 1 :| isNotParam "x" :| isNotZero "x"-    -- , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x"--    , "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isPositive "x"-    , (powabs "x" "y") * (powabs "x" "z") :=> powabs "x" ("y" + "x")-    , ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isPositive "x"+    -- B9: x^y * x^z = x^(y+z)+    , NAry EMul [Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'z')))]+        :==:+      Fixed (Bin Power (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')]))+        :| isPositive "x"+    -- B10: |x|^y * |x|^z = |x|^(y+z)  (fixed: target used "y+x" instead of "y+z")+    , NAry EMul [Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'z')))]+        :=>+      Fixed (Bin PowerAbs (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')]))+    -- B11: (x^y)^z = x^(y*z)+    , Fixed (Bin Power (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))) (VarPat 'z'))+        :==:+      Fixed (Bin Power (VarPat 'x') (NAry EMul [Ch (VarPat 'y'), Ch (VarPat 'z')]))+        :| isPositive "x"     , powabs (powabs "x" "y") "z" :=> powabs "x" ("y" * "z")-    , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isPositive "x" :| isPositive "y"--    --, "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isInteger "y" :| isInteger "z"  :| isNotZero "x"-    --, ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"-    --, ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isInteger "z" :| isNotZero "x" :| isNotZero "y"-     ]  rewritesWithConstant :: [Rule] rewritesWithConstant =     [-      "x" * "x" :=> "x" ** 2-    , "x" - "x" :=> 0+      "x" - "x" :=> 0     , "x" / "x" :=> 1 :| isNotZero "x"     , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isPositive "x"     , 1 ** "x" :=> 1     , powabs 1 "x" :=> 1     , log (sqrt "x") :=> 0.5 * log "x" :| isNotParam "x"-    , "x" ** (1/2)   :==: sqrt "x" -- <==>+    , "x" ** (1/2)   :==: sqrt "x"     , powabs "x" (1/2) :=> sqrt (abs "x")     , "x" ** (1/3) :==: Fixed (Uni Cbrt "x")-    , 0 * "x" :=> 0 :| isValid "x" -- :| isNotParam "x"+    -- B4: 0 * rest = 0 (provided every factor is valid)+    , NAry EMul [Ch (Fixed (Const 0)), Rest '1'] :=> 0 :| allValidRest '1'     , 0 ** "x" :=> 0 :| isPositive "x"     , powabs 0 "x" :=> 0-    , 0 - "x" :=> negate "x"-    , "x" + negate "y" :==: "x" - "y"+    -- n-ary cancellation: x + y - x = y+    , NAry EAdd [ Ch "a"+                , Ch (NAry EMul [ Ch (Fixed (Const (-1.0))), Ch "a" ])+                , Rest 'r' ]+        :=> NAry EAdd [Rest 'r']+    -- combining like terms: x + x = 2*x+    , NAry EAdd [ Ch "a", Ch "a", Rest 'r' ]+        :=> NAry EAdd [ Ch (2 * "a"), Rest 'r' ]     ] rewritesWithParam :: [Rule] rewritesWithParam =     [-    --  "x" * "x" :=> "x" ** Fixed (Param 0)       "x" - "x" :=> Fixed (Param 0)     , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"     , 1 ** "x" :=> Fixed (Param 0)     , powabs 1 "x" :=> Fixed (Param 0)-    -- , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"     ]  rewritesSimple :: [Rule]-rewritesSimple =-    [-      "x" * "y" :=> "y" * "x"-    , "x" + "y" :=> "y" + "x"-    , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"-    , ("x" + "y") + "z" :=> "x" + ("y" + "z")-    , ("x" * "y") * "z" :=> "x" * ("y" * "z")-    , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")-    , "x" - ("y" + "z") :=> ("x" - "y") - "z" -- TODO: check that I don't this-    , "x" - ("y" - "z") :=> ("x" - "y") + "z" -- TODO-    , ("x" * "y") / "z" :=> ("x" / "z") * "y" :| isNotZero "z" -- TODO: inv(x) <=> x^-1 , x/y <=> x*y^-1-    , "x" * ("y" / "z") :=> ("x" / "z") * "y" :| isNotZero "z" -- ^-    , "x" / ("y" * "z") :=> ("x" / "z") / "y" :| isNotZero "z" -- ^ TODO: 0 ^-1 check-    , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"-    , ("w" * "x") - ("z" * "x") :=> ("w" - "z") * "x" -- TODO: handle sub :| isConstPt "w" :| isConstPt "z"-    , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y")-    , log (exp "x")  :=> "x"-    , exp (log "x")  :=> "x"-    , log ("x" * "y") :=> log "x" + log "y"-    , log ("x" ** "y") :=> "y" * log "x"-    , abs ("x" * "y") :=> abs "x" * abs "y"-    , abs ("x" ** "y") :=> abs "x" ** "y"-    , abs ("x" - "y") :=> abs ("y" - "x")-    , recip (recip "x") :=> "x" :| isNotZero "x"-    , "x" * "x" :=> "x" ** Fixed (Param 0)-    , "x" - "x" :=> Fixed (Param 0)-    , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"-    , 1 ** "x" :=> Fixed (Param 0)-    , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"-    ]+rewritesSimple = rewriteBasic <> constReduction <> rewritesFun powabs l r = Fixed (Bin PowerAbs l r)  -- | default cost function for simplification@@ -267,14 +271,14 @@ rewritesParams :: [Rule] rewritesParams = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithParam --- | simplify using the default parameters +-- | simplify using the default parameters simplifyEqSatDefault :: Fix SRTree -> Fix SRTree-simplifyEqSatDefault t = eqSat t rewrites myCost 30 `evalState` emptyGraph+simplifyEqSatDefault t = eqSat t rewrites myCost 30 `evalState` emptyGraphNoTrack  -- | simplifies with custom parameters simplifyEqSat :: [Rule] -> CostFun -> Int -> Fix SRTree -> Fix SRTree simplifyEqSat rwrts costFun it t = eqSat t rwrts costFun it `evalState` emptyGraph  -- | apply a single step of merge-only using default rules-applyMergeOnlyDftl :: Monad m => CostFun -> EGraphST m ()+applyMergeOnlyDftl :: ClassStore m => CostFun -> EGraphST m () applyMergeOnlyDftl costFun = applySingleMergeOnlyEqSat costFun rewrites