packages feed

srtree-db 0.1.0.0 → 0.1.1.0

raw patch · 8 files changed

+602/−45 lines, 8 filesdep +filepathdep +optparse-applicativedep +randomdep ~srtreenew-component:exe:srtree-dbPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: filepath, optparse-applicative, random

Dependency ranges changed: srtree

API changes (from Hackage documentation)

- Algorithm.EqSat.Storage.Import: importEqs :: SqlBackend db => db -> String -> [(Fix SRTree, [Target], Maybe Double)] -> IO (Either String ImportSummary)
+ Algorithm.EqSat.Storage.Import: importEqs :: SqlBackend db => db -> Maybe String -> [(Fix SRTree, [Target], Maybe Double)] -> IO (Either String ImportSummary)

Files

ChangeLog.md view
@@ -1,5 +1,9 @@ # Changelog for srtree-db +## 0.1.1.0++- Added cli tools to populate and fit data into a database + ## 0.1.0.0  - Initial release
+ app/EqSat.hs view
@@ -0,0 +1,96 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards #-}++module EqSat+  ( EqSatOpts(..)+  , eqsatParser+  , runEqSatCmd+  ) where++import Control.Exception (bracket, SomeException, catch, displayException)+import Control.Monad.State.Strict (execStateT)+import qualified Data.IntMap.Strict as IntMap+import qualified Data.Text as T+import Options.Applicative++import Data.SRTree (SRTree(..))+import Algorithm.EqSat (runEqSat)+import Algorithm.EqSat.Egraph (EGraph(..), EClassPageStore(..))+import Algorithm.EqSat.Simplify (rewrites, rewritesParams, myCost)+import Algorithm.EqSat.Storage.Backend (SqlBackend)+import Algorithm.EqSat.Storage.SQLite (loadGraphLazy, saveGraph, flushStore)+import Algorithm.EqSat.Storage.Query (getOrCreateDataset)++import Database.SQLite3 (Database, open, close)++-- | CLI options for the eqsat sub-command.+data EqSatOpts = EqSatOpts+  { eqsatDb      :: String+  , eqsatDataset :: String+  , eqsatSteps   :: Int+  , eqsatRuleset :: String+  } deriving (Show)++eqsatParser :: Parser EqSatOpts+eqsatParser = EqSatOpts+  <$> strOption+      ( long "db"+      <> metavar "FILE"+      <> help "SQLite database file path" )+  <*> strOption+      ( long "dataset"+      <> metavar "NAME"+      <> help "Dataset name" )+  <*> option auto+      ( long "steps"+      <> value 1+      <> metavar "N"+      <> help "Number of eqsat iterations" )+  <*> strOption+      ( long "ruleset"+      <> value "default"+      <> metavar "RULESET"+      <> help "Rule set: default or params" )++-- | Run the eqsat sub-command.+runEqSatCmd :: EqSatOpts -> IO ()+runEqSatCmd EqSatOpts{..} = do+  let rules = case eqsatRuleset of+                "params" -> rewritesParams+                _        -> rewrites++  putStrLn $ "Loading paged graph from " ++ eqsatDb ++ "..."+  r <- withSQLite eqsatDb $ \db -> do+    dsid <- getOrCreateDataset db eqsatDataset+    er <- loadGraphLazy db dsid+    case er of+      Left err -> pure (Left err)+      Right eg -> do+        let classCount = IntMap.size (_eClass eg)+        putStrLn $ "Loaded " ++ show classCount ++ " e-classes"+        putStrLn $ "Running " ++ show eqsatSteps ++ " steps of eqsat with '"+                 ++ eqsatRuleset ++ "' rules..."+        let go g = execStateT (runEqSat myCost rules eqsatSteps) g+        eg' <- go eg+        let classCount' = IntMap.size (_eClass eg')+        flushStore eg'+        saveResult <- saveGraph db dsid eg'+        case saveResult of+          Left err -> pure (Left ("saveGraph failed: " ++ err))+          Right _  -> do+            -- clear frontier after full eqsat+            case _classStore eg' of+              Nothing -> pure ()+              Just h  -> cpsEndFrontier h+            pure (Right (classCount, classCount'))++  case r of+    Left err -> putStrLn $ "eqsat failed: " ++ err+    Right (before, after) -> do+      putStrLn $ "After eqsat: " ++ show after ++ " e-classes ("+               ++ show (after - before) ++ " change from " ++ show before ++ ")"+      putStrLn $ "Saved to " ++ eqsatDb ++ " [dataset: " ++ eqsatDataset ++ "]"++-- | Open a SQLite database, run an action, and close it.+withSQLite :: String -> (Database -> IO a) -> IO a+withSQLite path = bracket (open (T.pack path)) close
+ app/FitData.hs view
@@ -0,0 +1,205 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards #-}++module FitData+  ( FitDataOpts(..)+  , fitdataParser+  , runFitData+  ) where++import Control.Exception (bracket, SomeException, catch, displayException)+import Control.Monad.State.Strict (runStateT)+import Data.IORef+import Data.List (maximumBy)+import Data.Ord (comparing)+import qualified Data.Text as T+import qualified Data.Vector.Unboxed as VU+import Options.Applicative hiding (Const)+import System.Random (randomRIO)++import Data.SRTree (Fix(..), SRTree(..), Op(..), relabelParams, countParamsUniq, countNodes)+import Data.SRTree.Print (showExpr)+import Data.SRTree.Eval (Target)+import Data.SRTree.Datasets (loadDataset)+import Algorithm.SRTree.NonlinearOpt (minimizeNLL')+import Algorithm.SRTree.Likelihoods (Loss(..), Distribution(..), readLoss)+import Algorithm.SRTree.AD (ADBackEnd(..))+import Numeric.Optimization.NLOPT (LocalAlgorithm(..))+import Algorithm.EqSat.Egraph (EGraph(..), EClassId, getBestExpr, canonical)+import Algorithm.EqSat.Storage.Backend (SqlBackend(..), SqlValue(..), sqlToInt)+import Algorithm.EqSat.Storage.SQLite (loadGraphLazy, saveGraph, flushStore)+import Algorithm.EqSat.Storage.Query (getOrCreateDataset, writeDatasetFit)+import Algorithm.EqSat.Storage.Types (serializeTheta)+import Database.SQLite3 (Database, open, close)++-- | CLI options for the fitdata sub-command.+data FitDataOpts = FitDataOpts+  { fitdataDb        :: String+  , fitdataDataset   :: String+  , fitdataData      :: String+  , fitdataLoss      :: Loss+  , fitdataHasHeader :: Bool+  , fitdataNRep      :: Int+  , fitdataNIter     :: Int+  , fitdataBatchSize :: Int+  } deriving (Show)++fitdataParser :: Parser FitDataOpts+fitdataParser = FitDataOpts+  <$> strOption+      ( long "db"+      <> metavar "FILE"+      <> help "SQLite database file path" )+  <*> strOption+      ( long "dataset"+      <> metavar "NAME"+      <> help "Dataset name" )+  <*> strOption+      ( long "data"+      <> metavar "SPEC"+      <> help "Dataset CSV spec: file:start:end:target:features:yerr" )+  <*> option auto+      ( long "loss"+      <> value (NLL Gaussian)+      <> metavar "LOSS"+      <> help "Loss function (MSE, NLL Gaussian, etc.)" )+  <*> switch+      ( long "has-header"+      <> help "CSV has header row (default: True)" )+  <*> option auto+      ( long "n-rep"+      <> value 1+      <> metavar "N"+      <> help "Number of random restarts per expression" )+  <*> option auto+      ( long "n-iter"+      <> value 100+      <> metavar "N"+      <> help "Max NLopt iterations" )+  <*> option auto+      ( long "batch-size"+      <> value 100+      <> metavar "N"+      <> help "Fit N expressions per commit batch" )++-- | Run the fitdata sub-command.+runFitData :: FitDataOpts -> IO ()+runFitData opts = do+  let FitDataOpts{..} = opts+  -- 1. Load dataset+  putStrLn $ "Loading dataset: " ++ fitdataData+  ((xTrain, yTrain, _xVal, _yVal), (mYErr, _), _varnames, _target) <-+    loadDataset fitdataData fitdataHasHeader++  let nNoiseParams = case fitdataLoss of+        NLL Gaussian -> 1+        NLL ROXY     -> 3+        _            -> 0++  -- 2. Load paged graph and query for unfitted e-classes+  putStrLn $ "Loading paged graph from " ++ fitdataDb ++ "..."+  withSQLite fitdataDb $ \db -> do+    dsid <- getOrCreateDataset db fitdataDataset+    er <- loadGraphLazy db dsid+    case er of+      Left err -> putStrLn $ "loadGraphLazy failed: " ++ err+      Right eg -> do+        -- Query for unfitted e-classes+        unfitted <- queryUnfitted db dsid+        let total = length unfitted+        putStrLn $ "Found " ++ show total ++ " unfitted e-classes"++        if total == 0+          then putStrLn "Nothing to fit."+          else do+            -- Process in batches+            counter <- newIORef (0 :: Int)+            let processBatch [] = pure ()+                processBatch batch = do+                  mapM_ (fitOne db dsid eg xTrain yTrain mYErr fitdataLoss fitdataNIter fitdataNRep nNoiseParams counter) batch+                  -- Commit checkpoint+                  flushStore eg+                  putStrLn $ "  [checkpoint] committed batch"++            let batches = chunk fitdataBatchSize unfitted+            mapM_ processBatch batches++            fitted <- readIORef counter+            putStrLn $ "Fitted " ++ show fitted ++ "/" ++ show total+                     ++ " expressions"++-- | Fit a single e-class and write the result.+fitOne :: SqlBackend db+       => db -> Int -> EGraph+       -> [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double)+       -> Loss -> Int -> Int -> Int -> IORef Int -> EClassId -> IO ()+fitOne db dsid eg xTrain yTrain mYErr loss nIter nRep nNoiseParams counter eid = do+  n <- readIORef counter+  let label = "[" ++ show (n+1) ++ "]"+  -- Extract best expression from the e-class+  mTree <- runStateT (getBestExpr eid) eg >>= pure . Just . fst+  case mTree of+    Nothing -> do+      putStrLn $ label ++ " eclass " ++ show eid ++ ": could not extract expression (skipped)"+      modifyIORef' counter (+1)+    Just tree -> do+      let tree' = relabelParams tree+          np = countParamsUniq tree' + nNoiseParams+          sz = countNodes tree'+      -- Skip expressions with no parameters (pure constants/variables)+      if np == 0+        then do+          putStrLn $ label ++ " eclass " ++ show eid ++ ": no parameters (skipped)"+          modifyIORef' counter (+1)+        else do+          -- Try multiple random restarts, pick the best+          results <- mapM (\_ -> fitOneRandom xTrain yTrain mYErr loss nIter tree' np) [1..nRep]+          let (bestFitness, bestTheta) = maximumBy (comparing fst) results+          -- Write result to DB+          writeDatasetFit db dsid eid (Just bestFitness) Nothing+            (T.pack (serializeTheta [bestTheta])) sz+          putStrLn $ label ++ " eclass " ++ show eid ++ " (" ++ takeExpr tree' ++ "): fitness=" ++ showFit bestFitness+          modifyIORef' counter (+1)++-- | Fit an expression with a random initial theta.+fitOneRandom :: [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double)+             -> Loss -> Int -> Fix SRTree -> Int -> IO (Double, VU.Vector Double)+fitOneRandom xTrain yTrain mYErr loss nIter tree np = do+  theta0 <- VU.replicateM np (randomRIO (-1, 1))+  let (theta, lossVal, _) = minimizeNLL' VAR1 SingleThread loss mYErr nIter xTrain yTrain tree theta0+      fitness = negate lossVal+  pure (fitness, theta)++-- | Query for e-class IDs that are not yet fitted for a dataset.+queryUnfitted :: SqlBackend db => db -> Int -> IO [EClassId]+queryUnfitted db dsid = do+  rows <- queryDb db+    "SELECT e.eid FROM eclass e \+    \LEFT JOIN dataset_fit df ON df.eid = e.eid AND df.dataset_id = ? \+    \WHERE df.fitted IS NULL OR df.fitted = 0"+    [SqlInteger (fromIntegral dsid)]+  pure [ sqlToInt eid | [eid] <- rows ]++-- | Chunk a list into sub-lists of the given size.+chunk :: Int -> [a] -> [[a]]+chunk _ [] = []+chunk n xs = let (h, t) = splitAt n xs in h : chunk n t++-- | Take first few chars of an expression for display.+takeExpr :: Fix SRTree -> String+takeExpr t+  | length s > 40 = take 40 s ++ "..."+  | otherwise = s+  where s = showExpr t++-- | Format fitness for display.+showFit :: Double -> String+showFit f+  | f == (-1/0) = "-Infinity"+  | f == (1/0)  = "Infinity"+  | isNaN f     = "NaN"+  | otherwise   = show (fromIntegral (round (f * 1000) :: Int) / 1000 :: Double)++-- | Open a SQLite database, run an action, and close it.+withSQLite :: String -> (Database -> IO a) -> IO a+withSQLite path = bracket (open (T.pack path)) close
+ app/Ingest.hs view
@@ -0,0 +1,182 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards #-}++module Ingest+  ( IngestOpts(..)+  , ingestParser+  , runIngest+  ) where++import Control.Exception (bracket, SomeException, catch, displayException)+import Control.Monad (forM_, when, unless)+import Data.IORef+import qualified Data.ByteString.Char8 as B+import qualified Data.Text as T+import qualified Data.Vector.Unboxed as VU+import Options.Applicative+import System.IO (hIsEOF, hGetLine, stdin, openFile, IOMode(..), hClose, hPutStrLn, stderr, hFlush)+import System.Exit (exitFailure)+import System.Random (randomRIO)++import Data.SRTree (Fix(..), SRTree(..), relabelParams, countParamsUniq)+import Data.SRTree.Eval (Target)+import Data.SRTree.Datasets (loadDataset)+import Text.ParseSR (SRAlgs(..), parseSR)+import Algorithm.SRTree.NonlinearOpt (minimizeNLL')+import Algorithm.SRTree.Likelihoods (Loss(..), Distribution(..))+import Algorithm.SRTree.AD (ADBackEnd(..))+import Numeric.Optimization.NLOPT (LocalAlgorithm(..))+import Algorithm.EqSat.Storage.Import (importEqs, ImportSummary(..))+import Algorithm.EqSat.Storage.SQLite ()+import Database.SQLite3 (Database, open, close)++-- | CLI options for the ingest sub-command.+data IngestOpts = IngestOpts+  { ingestDb         :: String+  , ingestExprs      :: String+  , ingestDataset    :: String+  , ingestFormat     :: SRAlgs+  , ingestVarnames   :: String+  , ingestData       :: String+  , ingestFit        :: Bool+  , ingestLoss       :: Loss+  , ingestEqsatSteps :: Int+  , ingestReparam    :: Bool+  , ingestHasHeader  :: Bool+  } deriving (Show)++ingestParser :: Parser IngestOpts+ingestParser = IngestOpts+  <$> strOption+      ( long "db"+      <> metavar "FILE"+      <> help "SQLite database file path" )+  <*> strOption+      ( long "expressions"+      <> value ""+      <> metavar "FILE"+      <> help "File with one expression per line (empty = stdin)" )+  <*> strOption+      ( long "dataset"+      <> value ""+      <> metavar "NAME"+      <> help "Dataset name (required for fitting)" )+  <*> option auto+      ( long "format"+      <> value OPERON+      <> metavar "FORMAT"+      <> help "Expression format: OPERON, TIR, HL, BINGO, GOMEA, PYSR" )+  <*> strOption+      ( long "varnames"+      <> value "x0,x1,x2,x3,x4,x5"+      <> metavar "VARNAMES"+      <> help "Comma-separated variable names" )+  <*> strOption+      ( long "data"+      <> value ""+      <> metavar "SPEC"+      <> help "Dataset CSV spec: file:start:end:target:features:yerr" )+  <*> switch+      ( long "fit"+      <> help "Fit expressions after ingest" )+  <*> option auto+      ( long "loss"+      <> value (NLL Gaussian)+      <> metavar "LOSS"+      <> help "Loss function (MSE, NLL Gaussian, etc.)" )+  <*> option auto+      ( long "eqsat-steps"+      <> value 0+      <> metavar "N"+      <> help "Run N eqsat steps after ingest" )+  <*> switch+      ( long "reparam"+      <> help "Float constants to parameters" )+  <*> switch+      ( long "has-header"+      <> help "CSV has header row (default: True)" )++-- | Run the ingest sub-command.+runIngest :: IngestOpts -> IO ()+runIngest IngestOpts{..} = do+  -- Validate: --dataset is required when --fit is used+  when (ingestFit && null ingestDataset) $ do+    hPutStrLn stderr "Error: --dataset is required when --fit is used"+    exitFailure++  let mds = if null ingestDataset then Nothing else Just ingestDataset+      alg = ingestFormat+      varnames = ingestVarnames+      batchSize = 1000 :: Int++  -- Open the expression file (or stdin)+  h <- if null ingestExprs then pure stdin else openFile ingestExprs ReadMode++  -- Open DB+  putStrLn $ "Opening " ++ ingestDb ++ "..."+  db <- open (T.pack ingestDb)++  -- Process line by line, batch and insert+  putStrLn "Processing expressions..."+  totalRef   <- newIORef (0 :: Int)+  validRef   <- newIORef (0 :: Int)+  failedRef  <- newIORef (0 :: Int)+  classesRef <- newIORef (0 :: Int)+  batchRef   <- newIORef ([] :: [(Fix SRTree, [Target], Maybe Double)])++  let flushBatch = do+        batch <- readIORef batchRef+        unless (null batch) $ do+          r <- importEqs db mds (reverse batch)+          case r of+            Left err -> hPutStrLn stderr $ "  BATCH INSERT FAILED: " ++ err+            Right s  -> modifyIORef' classesRef (+ isClasses s)+          writeIORef batchRef []++      processLine line = do+        modifyIORef' totalRef (+1)+        if null line+          then pure ()+          else case parseSR alg (B.pack varnames) False (B.pack line) of+            Left err -> do+              modifyIORef' failedRef (+1)+              hPutStrLn stderr $ "  FAILED: " ++ line ++ " -- " ++ err+            Right tree -> do+              modifyIORef' validRef (+1)+              modifyIORef' batchRef ((relabelParams tree, [], Nothing) :)+              batch <- readIORef batchRef+              when (length batch >= batchSize) $ do+                flushBatch+                v <- readIORef validRef+                hPutStrLn stderr $ "  ... " ++ show v ++ " expressions processed"+                hFlush stderr++      loop = do+        done <- hIsEOF h+        if done then pure ()+        else do+          line <- hGetLine h+          processLine line+          loop++  loop+  flushBatch  -- insert any remaining expressions++  -- Close file handle+  unless (null ingestExprs) (hClose h)++  -- Summary+  total  <- readIORef totalRef+  valid  <- readIORef validRef+  failed <- readIORef failedRef+  classes <- readIORef classesRef+  putStrLn $ "Parsed " ++ show total ++ " expressions ("+           ++ show valid ++ " valid, " ++ show failed ++ " failed)"+  putStrLn $ "Imported into " ++ ingestDb+           ++ maybe "" (\d -> " [dataset: " ++ d ++ "]") mds+           ++ ": " ++ show classes ++ " e-classes"++  close db++  when (ingestEqsatSteps > 0) $+    putStrLn "(eqsat after ingest not yet implemented in standalone CLI)"
+ app/Main.hs view
@@ -0,0 +1,27 @@+{-# LANGUAGE OverloadedStrings #-}++module Main where++import Options.Applicative+import Ingest (IngestOpts, ingestParser, runIngest)+import EqSat  (EqSatOpts, eqsatParser, runEqSatCmd)+import FitData (FitDataOpts, fitdataParser, runFitData)++data Cmd = Ingest IngestOpts | EqSat EqSatOpts | FitData FitDataOpts++main :: IO ()+main = execParser cmdParser >>= dispatch++cmdParser :: ParserInfo Cmd+cmdParser = info (subcommands <**> helper) (progDesc "srtree-db: e-graph database CLI")+  where+    subcommands = subparser+      (  command "ingest"  (Ingest  <$> info (ingestParser <**> helper) (progDesc "Ingest expressions into DB"))+      <> command "eqsat"   (EqSat   <$> info (eqsatParser <**> helper) (progDesc "Run equality saturation"))+      <> command "fitdata" (FitData <$> info (fitdataParser <**> helper) (progDesc "Fit expressions to dataset"))+      )++dispatch :: Cmd -> IO ()+dispatch (Ingest  opts) = runIngest opts+dispatch (EqSat   opts) = runEqSatCmd opts+dispatch (FitData opts) = runFitData opts
src/Algorithm/EqSat/Storage/Import.hs view
@@ -68,14 +68,22 @@ -- structurally expanding every subexpression into its own e-class, then write -- the class pages in a final linear pass. Runs inside a single transaction -- (rolled back on error).-importEqs :: SqlBackend db => db -> String -> [(Fix SRTree, [Target], Maybe Double)] -> IO (Either String ImportSummary)-importEqs db ds eqs = do+--+-- When a dataset name is provided (@Just ds@), dataset_fit rows and+-- expression_index entries are written so fitness queries and dedup work.+-- When @Nothing@, only the structural e-graph is built (enode, eclass,+-- parent, cstore_page, meta) — the e-graph is reusable across datasets.+importEqs :: SqlBackend db => db -> Maybe String -> [(Fix SRTree, [Target], Maybe Double)] -> IO (Either String ImportSummary)+importEqs db mds eqs = do   createSchema db   -- content-address dedup queries by enode_key, but eclass_node's PK is   -- (eid, enode_key); index enode_key so those lookups are O(log n) not a scan.   execDb db "CREATE INDEX IF NOT EXISTS idx_eclass_node_enode_key ON eclass_node(enode_key)"-  dsid <- getOrCreateDataset db ds-  ref <- newIORef (ImportState 0)+  mdsid <- traverse (getOrCreateDataset db) mds+  -- Read the current next_id from the meta table so repeated importEqs calls+  -- don't collide on eclass eids.+  curNextId <- readMetaNextId db+  ref <- newIORef (ImportState curNextId)   r <- try $ do     execDb db "BEGIN"     -- stream the expression list through a fold (rather than forM_ + length@@ -83,18 +91,22 @@     -- is processed; the fold accumulator carries the count, so we never retain     -- the whole parsed expression list in memory.     n <- foldM (\c (t, theta, fit) -> do-                   (eid, h) <- insertTree ref db fit dsid t-                   writeDatasetFit db dsid eid fit Nothing (T.pack (serializeTheta theta)) h-                   -- record the root expression in the registry (keyed by its-                   -- canonical root node) so "was this expression seen/tested?"-                   -- can be answered per dataset.-                   mroot <- lookupClassNode db eid-                   forM_ mroot $ \en ->-                     runDb db-                       "INSERT OR REPLACE INTO expression_index (expression_key, eclass, dataset_id) VALUES (?, ?, ?)"-                       [ SqlText (T.pack (enodeKey en))-                       , SqlInteger (fromIntegral eid)-                       , SqlInteger (fromIntegral dsid) ]+                   (eid, h) <- insertTree ref db fit mdsid t+                   -- dataset_fit and expression_index only when a dataset is given+                   case mdsid of+                     Nothing -> pure ()+                     Just dsid -> do+                       writeDatasetFit db dsid eid fit Nothing (T.pack (serializeTheta theta)) h+                       -- record the root expression in the registry (keyed by its+                       -- canonical root node) so "was this expression seen/tested?"+                       -- can be answered per dataset.+                       mroot <- lookupClassNode db eid+                       forM_ mroot $ \en ->+                         runDb db+                           "INSERT OR REPLACE INTO expression_index (expression_key, eclass, dataset_id) VALUES (?, ?, ?)"+                           [ SqlText (T.pack (enodeKey en))+                           , SqlInteger (fromIntegral eid)+                           , SqlInteger (fromIntegral dsid) ]                    pure (c + 1)) 0 eqs     writeMeta db ref     writeAllPages db@@ -109,33 +121,33 @@       pure (Right (ImportSummary (stNextId st) (stNextId st) n))  -- | Insert a full tree bottom-up, returning its root e-class id and height.-insertTree :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Int -> Fix SRTree -> IO (EClassId, Int)-insertTree ref db fit dsid t = case unfix t of-  Var ix     -> insertNode ref db fit dsid (EVar ix) []-  Param ix   -> insertNode ref db fit dsid (EParam ix) []-  Const x    -> insertNode ref db fit dsid (EConst x) []+insertTree :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Maybe Int -> Fix SRTree -> IO (EClassId, Int)+insertTree ref db fit mdsid t = case unfix t of+  Var ix     -> insertNode ref db fit mdsid (EVar ix) []+  Param ix   -> insertNode ref db fit mdsid (EParam ix) []+  Const x    -> insertNode ref db fit mdsid (EConst x) []   Uni f sub  -> do-    (c, ch) <- insertTree ref db fit dsid sub-    insertNode ref db fit dsid (EUni f c) [(c, 1, ch)]-  Bin Add l r -> insertNAry ref db fit dsid EAdd l r-  Bin Mul l r -> insertNAry ref db fit dsid EMul l r+    (c, ch) <- insertTree ref db fit mdsid sub+    insertNode ref db fit mdsid (EUni f c) [(c, 1, ch)]+  Bin Add l r -> insertNAry ref db fit mdsid EAdd l r+  Bin Mul l r -> insertNAry ref db fit mdsid EMul l r   Bin op l r  -> do-    (lc, lh) <- insertTree ref db fit dsid l-    (rc, rh) <- insertTree ref db fit dsid r-    insertNode ref db fit dsid (EBin op lc rc) [(lc, 1, lh), (rc, 1, rh)]+    (lc, lh) <- insertTree ref db fit mdsid l+    (rc, rh) <- insertTree ref db fit mdsid r+    insertNode ref db fit mdsid (EBin op lc rc) [(lc, 1, lh), (rc, 1, rh)]  -- | Insert a flattened n-ary node (@Add@/@Mul@), merging nested same-op chains -- the same way 'mkENaryM' does (a child whose class holds a single same-op -- ENAry is flattened in). Children and their heights are read from the DB.-insertNAry :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Int -> NOp -> Fix SRTree -> Fix SRTree -> IO (EClassId, Int)-insertNAry ref db fit dsid op l r = do-  (c1, _) <- insertTree ref db fit dsid l-  (c2, _) <- insertTree ref db fit dsid r+insertNAry :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Maybe Int -> NOp -> Fix SRTree -> Fix SRTree -> IO (EClassId, Int)+insertNAry ref db fit mdsid op l r = do+  (c1, _) <- insertTree ref db fit mdsid l+  (c2, _) <- insertTree ref db fit mdsid r   flat <- flattenChildren db op [(c1, 1), (c2, 1)]   childsH <- forM (IntMap.toList flat) $ \(c, n) -> do     h <- classHeight db c     pure (c, n, h)-  insertNode ref db fit dsid (ENAry op flat) childsH+  insertNode ref db fit mdsid (ENAry op flat) childsH  -- | Flatten @n@ occurrences of @cid@ when its class holds exactly one ENAry of -- the same op (scaled by @n@); otherwise keep @cid@. Each child's node is read@@ -151,8 +163,8 @@  -- | Content-addressed insert of a single e-node. Returns the e-class id and -- height of the node (reusing the existing class when already present).-insertNode :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Int -> ENode -> [(EClassId, Int, Int)] -> IO (EClassId, Int)-insertNode ref db fit dsid en children = do+insertNode :: SqlBackend db => IORef ImportState -> db -> Maybe Double -> Maybe Int -> ENode -> [(EClassId, Int, Int)] -> IO (EClassId, Int)+insertNode ref db fit mdsid en children = do   let key   = enodeKey en       childs = dedupChildren children   mEid <- lookupEnodeId db key@@ -165,7 +177,7 @@       let eid = stNextId st           h   = 1 + maximum (0 : [ ch | (_, _, ch) <- childs ])       writeIORef ref (st { stNextId = eid + 1 })-      writeNode db eid en key childs h fit dsid+      writeNode db eid en key childs h fit mdsid       pure (eid, h)  -- | Merge duplicate child e-classes into multiplicities (e.g. @x0 - x0@ has@@ -178,8 +190,8 @@  -- | Write the relational rows for a brand-new e-node. Reverse parent edges go -- straight to the @parent@ table (reconstructed into pages by 'writeAllPages').-writeNode :: SqlBackend db => db -> EClassId -> ENode -> String -> [(EClassId, Int, Int)] -> Int -> Maybe Double -> Int -> IO ()-writeNode db eid en key children h fit dsid = do+writeNode :: SqlBackend db => db -> EClassId -> ENode -> String -> [(EClassId, Int, Int)] -> Int -> Maybe Double -> Maybe Int -> IO ()+writeNode db eid en key children h fit mdsid = do   runDb db "INSERT INTO enode (key, op, op_detail) VALUES (?, ?, ?)"     [ SqlText (T.pack key)     , SqlText (T.pack (enodeOpTag en))@@ -204,7 +216,9 @@       , SqlText (T.pack key) ]   -- every class gets a dataset_fit row so the graph is fitness-annotated per   -- dataset; the root's proper params/theta are written by the caller.-  writeDatasetFit db dsid eid fit Nothing "" h+  case mdsid of+    Nothing -> pure ()+    Just dsid -> writeDatasetFit db dsid eid fit Nothing "" h  -- | ENAry children as (class, multiplicity); empty for all other node shapes -- (their children live inline in the content key).@@ -306,7 +320,15 @@ writeMeta :: SqlBackend db => db -> IORef ImportState -> IO () writeMeta db ref = do   st <- readIORef ref-  runDb db "INSERT INTO meta (key, value) VALUES (?, ?)"+  runDb db "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)"     [ SqlText "next_id", SqlText (T.pack (show (stNextId st))) ]-  runDb db "INSERT INTO meta (key, value) VALUES (?, ?)"+  runDb db "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)"     [ SqlText "track_dbs", SqlText "1" ]++-- | Read the current next_id from the meta table (0 if no meta row exists).+readMetaNextId :: SqlBackend db => db -> IO Int+readMetaNextId db = do+  rows <- queryDb db "SELECT value FROM meta WHERE key = 'next_id'" []+  case rows of+    ([SqlText v] : _) -> pure (read (T.unpack v) :: Int)+    _                  -> pure 0
src/Algorithm/EqSat/Storage/SQLite.hs view
@@ -223,9 +223,9 @@  writeMeta :: SqlBackend db => db -> GraphRows -> IO () writeMeta db rows = do-  run db "INSERT INTO meta (key, value) VALUES (?, ?)"+  run db "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)"     [ SqlText "next_id", SqlText (T.pack (show (_grNextId rows))) ]-  run db "INSERT INTO meta (key, value) VALUES (?, ?)"+  run db "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)"     [ SqlText "track_dbs", SqlText (if _grTrackDBs rows then "1" else "0") ]  writeNodes :: SqlBackend db => db -> GraphRows -> IO ()
srtree-db.cabal view
@@ -1,6 +1,6 @@ cabal-version: 2.4 name:          srtree-db-version:       0.1.0.0+version:       0.1.1.0 synopsis:      SQL persistence and querying for srtree e-graphs description:   Reusable storage layer for srtree multiset e-graphs:                a driver-neutral serialization of the Algorithm.EqSat.Store@@ -43,7 +43,7 @@     , text >=1.2 && <2.2     , direct-sqlite >=2.3 && <2.4     , postgresql-libpq >=0.10 && <0.12-    , srtree >=3.0 && <3.1+    , srtree >=3.0.0.2 && <3.1     , vector >=0.12 && <0.14     , mtl >=2.2 && <2.4   default-language: Haskell2010@@ -64,6 +64,27 @@     , postgresql-libpq >=0.10 && <0.12     , srtree >=3.0 && <3.1     , srtree-db >=0.1 && <0.2+  default-language: Haskell2010++executable srtree-db+  hs-source-dirs:   app+  main-is:          Main.hs+  other-modules:    Ingest, EqSat, FitData+  ghc-options:      -O2 -threaded -rtsopts -with-rtsopts=-N+  build-depends:+      base >=4.14 && <5+    , optparse-applicative >=0.17 && <0.19+    , srtree-db >=0.1 && <0.2+    , srtree >=3.0 && <3.1+    , bytestring >=0.10 && <0.13+    , containers >=0.6 && <0.9+    , text >=1.2 && <2.2+    , vector >=0.12 && <0.14+    , directory >=1.3 && <1.4+    , filepath >=1.4 && <1.5+    , direct-sqlite >=2.3 && <2.4+    , random >=1.2 && <1.3+    , mtl >=2.2 && <2.4   default-language: Haskell2010  test-suite srtree-db-test