srtree-db-0.1.4.0: src/Algorithm/EqSat/Storage/SQLite.hs
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
-- | SQLite-backed persistence for srtree e-graphs.
--
-- 'saveGraph' serializes an in-memory e-graph into the normalized schema:
-- structure in @enode@/@enode_child@/@eclass@/@eclass_node@, per-class risk
-- metrics in @fit@, scalars in @meta@. 'loadGraph' reconstructs the e-graph
-- via 'Algorithm.EqSat.Store.importEGraph', recomputing parent pointers and
-- the derived range databases.
--
-- Cost / best / consts are NOT persisted: they are derived quantities, re-set
-- to defaults on load (queries, pattern matching and refitting do not need
-- them). ENAry children (eclass -> multiplicity) are stored in @enode_child@;
-- Uni/Bin children are embedded in the content key.
--
-- The module is written against 'Algorithm.EqSat.Storage.Backend' and is
-- driver-neutral: the same code drives the PostgreSQL backend
-- ('Algorithm.EqSat.Storage.Postgres').
module Algorithm.EqSat.Storage.SQLite
( saveGraph
, saveGraphStructure
, loadGraph
, loadGraphResident
, loadGraphResidentNoFit
, loadGraphResidentFor
, loadGraphLazy
, loadGraphLazyNoFit
, pushFit
, pushFitFit
, pushFitSplit
, refreshFitness
, query
, flushStore
, loadPagesBulk
, emptyPagedGraph
, createSchema
, createSchemaStructure
, createSchemaFit
) where
import Control.Monad (forM, forM_, when, foldM)
import Control.Exception (SomeException, catch, displayException)
import Control.Monad.Identity (runIdentity)
import Control.Monad.State.Strict (execStateT, runStateT)
import Data.Int (Int64)
import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import Data.List (foldl', intercalate)
import qualified Data.IntSet as IntSet
import qualified Data.IntMap as IntMap
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as Set
import qualified Data.Map.Strict as Map
import qualified Data.Set as RangeSet
import Data.Binary (decode, encode)
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as BS
import qualified Data.Text.Encoding as TE
import Database.SQLite3
( Database, SQLData(..), StepResult(..)
, bind, columns, exec, step, withStatement )
import Data.SRTree.Eval (Target)
import Algorithm.EqSat.Egraph
( EGraph(..), EGraphST, ClassStore, EClassId, ENode(..), Consts(..)
, EClassPageStore(..), EClass(..), EClassData(..)
, EGraphDB(..), emptyDB, canonical, insertRange, eChildren, eOpKey )
import Algorithm.EqSat.Build (populate)
import Algorithm.EqSat.Info (insertFitness)
import Algorithm.EqSat.Store
( GraphRows(..), EClassRow(..), exportEGraph, importEGraph, rebuildDBs )
import Algorithm.EqSat.Storage.Backend
( SqlValue(..), SqlBackend(..), sqlToInt, sqlToMaybeDouble, sqlToText, sqlToBlob )
import Algorithm.EqSat.Storage.ClassStore
( classStoreTable, openClassStore, allPages, classStoreHandle )
import Algorithm.EqSat.Storage.Query (readDatasetFit, writeDatasetFit, firstDatasetId)
import Algorithm.EqSat.Storage.Stream (streamRootsByOp)
import Algorithm.EqSat.Storage.Types
import Algorithm.EqSat.Storage.Schema (egraphSchemaSQL, structureSchemaSQL, fitSchemaSQL, createSchema, createSchemaStructure, createSchemaFit)
-- | Default cache capacity (pages) for the lazily paged e-class store.
-- Kept small: the cache is an LRU of deserialized EClass objects (each with
-- HashSets of nodes/parents), so a large cap wastes resident memory. The
-- out-of-core matcher streams roots directly from the DB and the extraction
-- path reads one class at a time, so a small working set suffices.
defaultClassCap :: Int
defaultClassCap = 1000
-- ---------------------------------------------------------------------------
-- SQLite driver instance
instance SqlBackend Database where
execDb = exec
runDb db sql params = withStatement db sql $ \stmt -> do
bind stmt (map toSqlData params)
_ <- step stmt
pure ()
insertIgnore db tail params = withStatement db ("INSERT OR IGNORE INTO " <> tail) $ \stmt -> do
bind stmt (map toSqlData params)
_ <- step stmt
pure ()
queryDb db sql params = withStatement db sql $ \stmt -> do
bind stmt (map toSqlData params)
go stmt []
where
go stmt acc = do
r <- step stmt
case r of
Done -> pure (reverse acc)
Row -> do
cols <- columns stmt
go stmt (map fromSqlData cols : acc)
foldQueryDb db sql params seed0 go = withStatement db sql $ \stmt -> do
bind stmt (map toSqlData params)
goRows seed0 stmt
where
goRows !acc stmt = do
r <- step stmt
case r of
Done -> pure acc
Row -> do
cols <- columns stmt
acc' <- go acc (map fromSqlData cols)
goRows acc' stmt
-- O(1)-memory candidate-root enumeration: stream the distinct e-class ids by
-- operator through a cursor, skipping the already-attempted set, and stop
-- after @budget@ rows, so the matcher never materializes the whole (operator
-- -> root set) index in RAM.
streamByOp db opDetail budget exclude = streamRootsByOp db opDetail budget exclude
-- Stream every page of a key-value table through a cursor, so a full pass
-- (e.g. pushFit) stays O(1) memory instead of materializing all pages.
streamPages db tbl k = withStatement db ("SELECT key, blob FROM " <> tbl) $ \stmt -> do
go stmt
where
go stmt = do
r <- step stmt
case r of
Done -> pure ()
Row -> do
cols <- columns stmt
let eid = case cols of (SQLInteger i : _) -> i; _ -> 0
blob = case cols of (_ : SQLBlob bs : _) -> bs
(_ : SQLText t : _) -> TE.encodeUtf8 t
_ -> BS.empty
k eid blob
go stmt
createSchemaDb db = mapM_ (exec db) egraphSchemaSQL
createSchemaDbStructure db = mapM_ (exec db) structureSchemaSQL
createSchemaDbFit db = mapM_ (exec db) fitSchemaSQL
toSqlData :: SqlValue -> SQLData
toSqlData (SqlInteger n) = SQLInteger n
toSqlData (SqlFloat d) = SQLFloat d
toSqlData (SqlText t) = SQLText t
toSqlData (SqlBlob bs) = SQLBlob bs
toSqlData SqlNull = SQLNull
fromSqlData :: SQLData -> SqlValue
fromSqlData (SQLInteger n) = SqlInteger n
fromSqlData (SQLFloat d) = SqlFloat d
fromSqlData (SQLText t) = SqlText t
fromSqlData (SQLBlob bs) = SqlBlob bs
fromSqlData SQLNull = SqlNull
-- | Driver-neutral parameterized query (abstracts the concrete backend).
query :: SqlBackend db => db -> Text -> [SqlValue] -> IO [[SqlValue]]
query = queryDb
-- | Driver-neutral parameterized statement (abstracts the concrete backend).
run :: SqlBackend db => db -> Text -> [SqlValue] -> IO ()
run = runDb
-- ---------------------------------------------------------------------------
-- writing
-- | Persist the full e-graph (structure + risk metrics + e-class pages).
-- Replaces any previously stored graph in this database.
--
-- Every canonical e-class is written as a serialized page to @cstore_page@ in
-- addition to the normalized relational schema, so the graph round-trips
-- through the lazily paged 'loadGraph' as well as the relational path.
--
-- On a paged graph the page store and the write-through-maintained relational
-- tables are already the authoritative live graph ('cpsInsert' writes every
-- class body, 'cpsRecordNode'/'cpsRecordCanonical' keep the structure tables
-- current), so 'saveGraph' only refreshes the @meta@ scalars. This avoids
-- materializing every page ('cpsAll') in RAM -- the O(n) spike that dominated
-- out-of-core persistence -- and never rebuilds the canonical/node tables from
-- the (bounded, partial) resident caches.
saveGraph :: SqlBackend db => db -> Int -> EGraph -> IO (Either String ())
saveGraph db dsid eg = saveGraphWith True (createSchema db) db (Just dsid) eg
-- | Save only the e-graph structure. No @dataset@/@dataset_fit@ rows are
-- written, and the serialized class pages have their dataset-specific fields
-- (fitness/DL/theta) cleared, so a structure-only database contains no
-- per-dataset data at rest.
saveGraphStructure :: SqlBackend db => db -> EGraph -> IO (Either String ())
saveGraphStructure db eg = saveGraphWith False (createSchemaStructure db) db Nothing eg
saveGraphWith
:: SqlBackend db
=> Bool -- ^ write dataset_fit rows too (legacy combined mode)
-> IO () -- ^ schema creation action
-> db -> Maybe Int -> EGraph
-> IO (Either String ())
saveGraphWith writeFit mkSchema db mdsid eg = do
mkSchema
let rows0 = exportEGraph eg
case _classStore eg of
Just _ -> commitWith $ do
writeMeta db rows0
Nothing -> do
gr <- graphClassRows eg
let rowsBase = rows0 { _grEClasses = gr }
rows = if writeFit
then rowsBase
else rowsBase { _grEClasses = IntMap.map stripRowDatasetInfo gr }
commitWith $ do
clearTables db
writeMeta db rows
writeNodes db rows
writeClasses db rows
case (writeFit, mdsid) of
(True, Just dsid) -> writeDatasetFitRows db dsid rows
_ -> pure ()
writeClassPages db rows
where
commitWith writes = do
execDb db "BEGIN"
result <- (writes >> execDb db "COMMIT" >> pure (Right ()))
`catch` \(e :: SomeException) -> do
-- roll back so a partial write never leaves the connection mid-transaction
execDb db "ROLLBACK"
pure (Left ("saveGraph failed: " <> displayException e))
pure result
-- | Clear per-dataset fields from class metadata. @_size@ is structural and is
-- deliberately kept.
stripDatasetInfo :: EClassData -> EClassData
stripDatasetInfo info = info { _fitness = Nothing, _dl = Nothing, _theta = [] }
stripRowDatasetInfo :: EClassRow -> EClassRow
stripRowDatasetInfo r = r { _rcInfo = stripDatasetInfo (_rcInfo r) }
-- | Enumerate every canonical e-class row of a graph. For a paged graph the
-- resident @_eClass@ is a bounded cache that may hold classes created or mutated
-- after the last flush, so it is the authoritative source for any class it
-- contains; the persisted pages supply the remainder (classes evicted from the
-- resident cache or never touched). The two are unioned with the resident rows
-- taking precedence, so edits made through either the IO or the pure instances
-- are never lost on 'saveGraph'. A fully resident graph reads the complete map
-- through 'exportEGraph'.
graphClassRows :: EGraph -> IO (IntMap.IntMap EClassRow)
graphClassRows eg = case _classStore eg of
Nothing -> pure (_grEClasses (exportEGraph eg))
Just h -> do
pages <- cpsAll h
let storeRows = IntMap.fromList [ mkRow ec | ec <- pages ]
resident = _grEClasses (exportEGraph eg)
pure (resident `IntMap.union` storeRows)
where
mkRow ec = (_eClassId ec, EClassRow (_eNodes ec) (_parents ec) (_height ec) (_info ec))
clearTables :: SqlBackend db => db -> IO ()
clearTables db = do
execDb db "DELETE FROM meta"
execDb db "DELETE FROM enode_child"
execDb db "DELETE FROM eclass_node"
execDb db "DELETE FROM enode"
execDb db "DELETE FROM eclass"
execDb db ("DELETE FROM " <> classStoreTable)
-- | Serialize and store every canonical e-class as a page in the page store
-- table, inside the calling transaction (no nested BEGIN/COMMIT).
writeClassPages :: SqlBackend db => db -> GraphRows -> IO ()
writeClassPages db rows =
forM_ (IntMap.toAscList (_grEClasses rows)) $ \(eid, r) ->
run db ("INSERT INTO " <> classStoreTable <> " (key, blob) VALUES (?, ?)")
[ SqlInteger (fromIntegral eid)
, SqlBlob (BL.toStrict (encode (EClass eid (_rcNodes r) (_rcParents r) (_rcHeight r) (_rcInfo r)))) ]
writeMeta :: SqlBackend db => db -> GraphRows -> IO ()
writeMeta db rows = do
run db "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)"
[ SqlText "next_id", SqlText (T.pack (show (_grNextId rows))) ]
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 ()
writeNodes db rows =
forM_ (HashMap.toList (_grENodeToEClass rows)) $ \(en, eid) -> do
let key = enodeKey en
-- Content-addressed inserts: distinct e-nodes can share a key (every NaN
-- constant serializes to "Const NaN" while Eq Double treats NaNs as
-- distinct, so the node map may hold several of them). The key is the
-- identity, so colliding rows are canonicalized with OR IGNORE instead of
-- aborting the whole save with a UNIQUE-constraint failure.
run db "INSERT OR IGNORE INTO enode (key, op, op_detail) VALUES (?, ?, ?)"
[ SqlText (T.pack key)
, SqlText (T.pack (enodeOpTag en))
, SqlText (T.pack (enodeOpDetail en)) ]
run db "INSERT OR IGNORE INTO eclass_node (eid, enode_key) VALUES (?, ?)"
[ SqlInteger (fromIntegral eid), SqlText (T.pack key) ]
forM_ (naryChildren en) $ \(c, n) ->
run db "INSERT OR IGNORE INTO enode_child (enode_key, child_eid, cnt) VALUES (?, ?, ?)"
[ SqlText (T.pack key)
, SqlInteger (fromIntegral c)
, SqlInteger (fromIntegral n) ]
-- | Children of an ENAry node as (class, multiplicity); empty otherwise.
naryChildren :: ENode -> [(EClassId, Int)]
naryChildren (ENAry _ m) = IntMap.toList m
naryChildren _ = []
writeClasses :: SqlBackend db => db -> GraphRows -> IO ()
writeClasses db rows =
forM_ (IntMap.toAscList (_grCanonical rows)) $ \(eid, canon) ->
run db "INSERT INTO eclass (eid, canonical, height) VALUES (?, ?, ?)"
[ SqlInteger (fromIntegral eid)
, SqlInteger (fromIntegral canon)
, SqlInteger (fromIntegral (maybe 0 _rcHeight (IntMap.lookup eid (_grEClasses rows)))) ]
-- | Write the per-(dataset, e-class) fitness rows for every class in a graph.
writeDatasetFitRows :: SqlBackend db => db -> Int -> GraphRows -> IO ()
writeDatasetFitRows db dsid rows =
forM_ (IntMap.toAscList (_grEClasses rows)) $ \(eid, r) -> do
let info = _rcInfo r
writeDatasetFit db dsid eid (_fitness info) (_dl info)
(T.pack (serializeTheta (_theta info))) (_size info)
-- ---------------------------------------------------------------------------
-- reading
-- | Read a dataset's fit rows into an eid-keyed map. 'Nothing' means "no
-- dataset" (structure-only load): no dataset tables are queried, so this works
-- on a pure structure database.
readFitFor
:: SqlBackend db
=> Maybe Int -> db
-> IO (IntMap.IntMap (Maybe Double, Maybe Double, Int, [Target]))
readFitFor Nothing _ = pure IntMap.empty
readFitFor (Just ds) db = do
rows <- readDatasetFit db ds
pure $ IntMap.fromList
[ (eid, (f, d, sz, parseTheta (T.unpack th)))
| (eid, (f, d, sz, th)) <- rows ]
-- | Apply a dataset's fit rows to decoded e-classes.
--
-- @clearMissing@ controls classes that have no row in the selected dataset:
-- * 'True' - clear fitness/DL/theta (split mode: fitness is dataset state
-- and a class that was not fitted for this dataset is
-- unevaluated).
-- * 'False' - keep the values embedded in the page blob (legacy combined
-- databases, where paged writes historically stored fitness in
-- the page and not every class had a @dataset_fit@ row).
--
-- @_size@ is structural; a positive size from the fit row is accepted for
-- compatibility with older fit writers, otherwise the structural size is kept.
applyFitMap
:: Bool
-> IntMap.IntMap (Maybe Double, Maybe Double, Int, [Target])
-> IntMap.IntMap EClass
-> IntMap.IntMap EClass
applyFitMap clearMissing fitMap = IntMap.mapWithKey apply
where
apply eid ec =
let info0 = _info ec
info1 = if clearMissing then stripDatasetInfo info0 else info0
ec0 = ec { _info = info1 }
in case IntMap.lookup eid fitMap of
Nothing -> ec0
Just (f, d, s, th) ->
ec0 { _info = info1
{ _fitness = f
, _dl = d
, _size = if s > 0 then s else _size info1
, _theta = th
} }
-- | Reconstruct the e-graph stored by a previous 'saveGraph'.
--
-- When the database carries e-class pages (@cstore_page@), the graph is
-- restored from those pages with an 'EClassPageStore' handle installed (so
-- subsequent mutations are written through to the store) and the relational
-- tables supply structure (canonical map, node -> class) and the @fit@ table
-- supplies the current risk metrics. Databases written without pages fall
-- back to the fully relational path.
loadGraph :: SqlBackend db => db -> IO (Either String EGraph)
loadGraph db = do
m <- readMeta db
case m of
Nothing -> pure (Left "srtree-db: no e-graph stored in this database")
Just (nextId, trackDBs) -> do
enodes <- readNodes db
ecLst <- readClasses db
-- risk metrics come from @dataset_fit@ (the legacy @fit@ table is gone);
-- the legacy non-dataset loader uses the first dataset's values.
mdsid <- firstDatasetId db
fitMap <- readFitFor mdsid db
let canon = IntMap.fromList [ (eid, c) | (eid, c, _) <- ecLst ]
nodeToEClass = HashMap.fromList enodes
ps <- openClassStore db defaultClassCap 1000
pages <- allPages ps
if null pages
then do
-- fully relational path (databases written before the page store)
-- Parent pointers are recomputed from nodeToEClass (the parent table
-- has been removed; buildClasses handles the fallback).
let classes = buildClasses canon nodeToEClass IntMap.empty fitMap (IntMap.fromList [ (eid, h) | (eid, _, h) <- ecLst ])
rows = GraphRows canon nodeToEClass classes nextId trackDBs
pure (importEGraph rows)
else do
-- paged path: classes come from the serialized pages; the @fit@
-- table overrides their risk metrics (it is the current DB truth).
let classes = applyFitMap False fitMap
(IntMap.fromList [ (eid, decode (BL.fromStrict page)) | (eid, page) <- pages ])
toRow eid ec = EClassRow (_eNodes ec) (_parents ec) (_height ec) (_info ec)
rows = GraphRows canon nodeToEClass (IntMap.mapWithKey toRow classes) nextId trackDBs
case importEGraph rows of
Left err -> pure (Left err)
Right eg -> pure (Right eg { _classStore = Just (classStoreHandle ps) })
-- | Load a fully resident graph (no paged store). All classes are in the
-- resident map, the pattern trie is built, and no I/O occurs during eqsat.
-- Used for benchmarking in-memory vs paged performance.
-- Note: if some eclasses are referenced by nodes but don't have pages
-- (e.g. created during eqsat but not flushed), this will skip them.
loadGraphResident :: SqlBackend db => db -> IO (Either String EGraph)
loadGraphResident db = firstDatasetId db >>= \mdsid -> loadGraphResidentWithFit False mdsid db
-- | Load a structure-only resident graph. No dataset tables are queried, so
-- every class starts unevaluated. Safe on a pure structure database.
loadGraphResidentNoFit :: SqlBackend db => db -> IO (Either String EGraph)
loadGraphResidentNoFit db = loadGraphResidentWithFit True Nothing db
-- | Load a resident graph applying fit rows for a specific dataset in the same
-- database.
loadGraphResidentFor :: SqlBackend db => db -> Int -> IO (Either String EGraph)
loadGraphResidentFor db dsid = loadGraphResidentWithFit True (Just dsid) db
loadGraphResidentWithFit :: SqlBackend db => Bool -> Maybe Int -> db -> IO (Either String EGraph)
loadGraphResidentWithFit clearMissing mdsid db = do
m <- readMeta db `catch` \(_ :: SomeException) -> pure Nothing
case m of
Nothing -> pure (Left "srtree-db: no e-graph stored in this database")
Just (nextId, trackDBs) -> do
enodes <- readNodes db
ecLst <- readClasses db
fitMap <- readFitFor mdsid db
let canon = IntMap.fromList [ (eid, c) | (eid, c, _) <- ecLst ]
nodeToEClass = HashMap.fromList enodes
ps <- openClassStore db defaultClassCap 1000
pages <- allPages ps
if null pages
then do
let classes = buildClasses canon nodeToEClass IntMap.empty fitMap (IntMap.fromList [ (eid, h) | (eid, _, h) <- ecLst ])
rows = GraphRows canon nodeToEClass classes nextId trackDBs
pure (importEGraph rows)
else do
let classes = applyFitMap clearMissing fitMap
(IntMap.fromList [ (eid, decode (BL.fromStrict page)) | (eid, page) <- pages ])
-- Filter nodeToEClass to only reference eclasses that have pages
validEids = IntMap.keysSet classes
nodeToEClass' = HashMap.filter (`IntSet.member` validEids) nodeToEClass
toRow eid ec = EClassRow (_eNodes ec) (_parents ec) (_height ec) (_info ec)
rows = GraphRows canon nodeToEClass' (IntMap.mapWithKey toRow classes) nextId trackDBs
-- Note: _classStore = Nothing, so the Identity/State instance is used
-- (no I/O during eqsat, pure IntMap/HashMap lookups)
pure (importEGraph rows)
-- | Write back any pending dirty e-class pages when the graph carries a
-- paged store (a no-op on a fully resident graph). Call this at durable
-- commit points (e.g. rewrite-loop iteration boundaries).
flushStore :: EGraph -> IO ()
flushStore eg = case _classStore eg of
Nothing -> pure ()
Just h -> cpsFlush h
-- | Seed the derived DBs (pattern trie and size/fitness/DL range DBs) purely
-- from the structure and @fit@ tables, without materializing any e-class page.
-- This is the lazy path's analogue of 'rebuildDBs' (which reads @_eClass@) for
-- a graph whose resident e-class map starts empty ('loadGraphLazy').
--
-- Pattern-match entries are structural only: unlike the eager path we do not
-- substitute known-constant classes (that would require reading the pages).
-- Range/size/unevaluated sets come from the @fit@ table, which is sourced from
-- the DB just like in 'rebuildDBs'.
seedEDB
:: Int -> Bool
-> HashMap.HashMap ENode EClassId
-> IntMap.IntMap (Maybe Double, Maybe Double, Int, [Target])
-> EGraphDB
seedEDB nextId trackDBs nodeToEClass fitMap =
IntMap.foldlWithKey' step pat fitMap
where
trie0 :: EGraphDB
trie0 = (emptyDB){ _nextId = nextId, _trackDBs = trackDBs }
-- pattern trie: one path per (root class, children) per operator
pat = HashMap.foldlWithKey' addNode trie0 nodeToEClass
addNode db en eid =
let ids = eid : eChildren en
op = eOpKey en
cur = Map.lookup op (_patDB db)
in case populate cur ids of
Nothing -> db
Just t -> db { _patDB = Map.insert op t (_patDB db) }
-- size/fitness/DL range DBs + unevaluated set
step db eid (fitM, dlM, sz, _theta) =
let db1 = db { _sizeDB = IntMap.insertWith IntSet.union sz (IntSet.singleton eid) (_sizeDB db) }
db2 = case fitM of
Nothing -> db1 { _unevaluated = IntSet.insert eid (_unevaluated db1) }
Just fn -> db1 { _fitRangeDB = insertRange eid fn (_fitRangeDB db1)
, _sizeFitDB = IntMap.insertWith RangeSet.union sz (RangeSet.singleton (fn, eid)) (_sizeFitDB db1) }
db3 = case dlM of
Nothing -> db2
Just dn -> db2 { _dlRangeDB = insertRange eid dn (_dlRangeDB db2)
, _sizeDLDB = IntMap.insertWith RangeSet.union sz (RangeSet.singleton (dn, eid)) (_sizeDLDB db2) }
in db3
-- | Minimal DB seed for the lazily paged path: the base scalars (next id,
-- tracking) with NO pattern trie and NO size/fitness/DL range DBs. The
-- out-of-core eqsat streams both matcher paths ('matchStreamCached' and
-- 'matchNAryWith') directly from the backing store via 'streamRoots', so the
-- O(nodes) in-RAM @_patDB@ trie is not built at all here. ('_nodeToEClass' is
-- kept only to populate the EGraph's node->class map.)
seedEDBPaged
:: Int -> Bool
-> HashMap.HashMap ENode EClassId
-> Int -- ^ resident class cache capacity
-> Int -- ^ node-to-class cache capacity
-> Int -- ^ canonical map cache capacity
-> EGraphDB
seedEDBPaged nextId trackDBs _nodeToEClass residentCap nodeCap canonicalCap =
(emptyDB){ _nextId = nextId, _trackDBs = trackDBs
, _residentCap = residentCap, _nodeCap = nodeCap, _canonicalCap = canonicalCap }
-- | Reconstruct an e-graph for out-of-core use: like 'loadGraph' but the
-- resident e-class map is left empty and an 'EClassPageStore' handle is
-- installed so classes are streamed in and out of a bounded cache. Structure
-- (canonical map, node -> class) and the derived DBs come from the relational
-- tables; individual classes are fetched lazily from the page store.
--
-- This bounds peak memory (the whole class set is never resident at once).
-- Use it with 'MonadIO'-based ('ClassStore') operations; the pure instances
-- expect a complete resident map and are not suitable for a lazy graph.
--
-- On the paged path the resident @_canonicalMap@/@_eNodeToEClass@ start EMPTY
-- (bounded caches): canonical/node lookups fall back to the live relational
-- tables ('cpsCanonicalOf'/'cpsNodeToClass'), which the write-through keeps
-- current, so nothing O(nodes) is materialized at load.
loadGraphLazy :: SqlBackend db => db -> Int -> Int -> Int -> Int -> IO (Either String EGraph)
-- db dsid residentCap nodeCap canonicalCap
loadGraphLazy db dsid = loadGraphLazyWithFit False (Just dsid) db
-- | Structure-only lazy load: no dataset tables are queried, so every class
-- starts unevaluated. Safe on a pure structure database.
loadGraphLazyNoFit :: SqlBackend db => db -> Int -> Int -> Int -> IO (Either String EGraph)
loadGraphLazyNoFit db = loadGraphLazyWithFit True Nothing db
loadGraphLazyWithFit
:: SqlBackend db
=> Bool -> Maybe Int -> db -> Int -> Int -> Int -> IO (Either String EGraph)
loadGraphLazyWithFit clearMissing mdsid db residentCap nodeCap canonicalCap = do
m <- readMeta db `catch` \(_ :: SomeException) -> do
-- The egraph tables may not exist yet on a fresh database (e.g. first
-- insert); treat that as "no e-graph stored" rather than crashing.
pure Nothing
case m of
Nothing -> pure (Left "srtree-db: no e-graph stored in this database")
Just (nextId, trackDBs) -> do
fitMap <- readFitFor mdsid db
let -- slim (no theta) map for the out-of-core path: the eqsat matcher's
-- conditions read only class _consts, never theta, so attaching it on
-- every page read would retain O(#classes) target vectors in RAM.
fitSlim = IntMap.map (\(f, d, sz, _) -> (f, d, sz)) fitMap
ps <- openClassStore db defaultClassCap 1000
hasPages <- storeHasPages db
if not hasPages
then do
-- fully relational fallback (databases written before the page store)
enodes <- readNodes db
ecLst <- readClasses db
let canon0 = IntMap.fromList [ (eid, c) | (eid, c, _) <- ecLst ]
rep eid = IntMap.findWithDefault eid eid canon0
nodeToEClass0 = HashMap.fromList enodes
nodeToEClass = HashMap.map rep nodeToEClass0
classes = buildClasses canon0 nodeToEClass IntMap.empty fitMap
(IntMap.fromList [ (eid, h) | (eid, _, h) <- ecLst ])
rows = GraphRows canon0 nodeToEClass classes nextId trackDBs
pure (importEGraph rows)
else do
-- lazily paged: empty resident maps (canonical / node -> class start
-- as bounded caches backed by the live relational tables), a store
-- handle, and no pattern trie / range DBs. Fitness is dataset
-- metadata, applied on page reads (not baked into the pages).
let base = classStoreHandle ps
h = base { cpsLookup = \eid ->
fmap (fmap (applyDsFit clearMissing fitSlim eid)) (cpsLookup base eid) }
eDB = seedEDBPaged nextId trackDBs HashMap.empty residentCap nodeCap canonicalCap
eg = EGraph IntMap.empty HashMap.empty IntMap.empty eDB (Just h)
pure (Right eg)
-- | An empty out-of-core paged graph (no classes yet) with the given cache
-- capacities. Used to seed a fresh database on the first 'DBInsert', so
-- inserting into an empty DB works instead of failing with "no e-graph stored".
emptyPagedGraph :: SqlBackend db => db -> Int -> Int -> Int -> IO EGraph
emptyPagedGraph db residentCap nodeCap canonicalCap = do
ps <- openClassStore db defaultClassCap 1000
let eDB = seedEDBPaged 0 True HashMap.empty residentCap nodeCap canonicalCap
eg = EGraph IntMap.empty HashMap.empty IntMap.empty eDB (Just (classStoreHandle ps))
pure eg
-- | Apply a dataset's fitness metadata to a class read from the structural page
-- store (fitness/dl/size are dataset-specific, so they are attached on read
-- rather than stored in the page blob). Theta is deliberately NOT attached on
-- the out-of-core path: the eqsat matcher's conditions read only @_consts@, so
-- keeping the per-class @[Target]@ vectors out of the resident read path avoids
-- O(#classes) memory; theta stays in @dataset_fit@ for the query path.
applyDsFit
:: Bool
-> IntMap.IntMap (Maybe Double, Maybe Double, Int)
-> EClassId -> EClass -> EClass
applyDsFit clearMissing m eid ec =
let info0 = _info ec
info1 = if clearMissing then stripDatasetInfo info0 else info0
ec0 = ec { _info = info1 }
in case IntMap.lookup eid m of
Nothing -> ec0
Just (f, dl, sz) ->
ec0 { _info = info1
{ _fitness = f
, _dl = dl
, _size = if sz > 0 then sz else _size info1
} }
-- | Cheap emptiness test for the page store (avoids materializing every page
-- blob just to pick the relational vs paged load path).
storeHasPages :: SqlBackend db => db -> IO Bool
storeHasPages db = do
rows <- query db ("SELECT 1 FROM " <> classStoreTable <> " LIMIT 1") []
pure (not (null rows))
readMeta :: SqlBackend db => db -> IO (Maybe (Int, Bool))
readMeta db = do
rows <- query db "SELECT key, value FROM meta" []
let m = HashMap.fromList [ (sqlToText k, sqlToText v) | [k, v] <- rows ]
case HashMap.lookup "next_id" m of
Nothing -> pure Nothing
Just v -> pure (Just (fromMaybe 0 (listToMaybe [ i | (i, "") <- reads (T.unpack v) ])
, HashMap.lookupDefault "0" "track_dbs" m == "1"))
-- | Read (e-node, its e-class) pairs from the enode + eclass_node tables.
readNodes :: SqlBackend db => db -> IO [(ENode, EClassId)]
readNodes db = do
rows <- query db
"SELECT n.enode_key, n.eid FROM eclass_node n JOIN enode e ON e.key = n.enode_key" []
pure (catMaybes
[ do
en <- parseEnodeKey (T.unpack (sqlToText k))
pure (en, sqlToInt eid)
| [k, eid] <- rows ])
-- | Read (eid, canonical, height) triples.
readClasses :: SqlBackend db => db -> IO [(EClassId, EClassId, Int)]
readClasses db = do
rows <- query db "SELECT eid, canonical, height FROM eclass" []
pure [ (sqlToInt eid, sqlToInt c, sqlToInt h) | [eid, c, h] <- rows ]
-- | Rebuild @_grEClasses@ rows: only canonical roots carry real class rows.
-- Parent pointers are recomputed from the node -> class map (the parent table
-- has been removed; the cstore_page blobs store the authoritative parent set).
buildClasses
:: IntMap.IntMap EClassId -- ^ canonical eid -> eid (self-map for roots)
-> HashMap.HashMap ENode EClassId -- ^ node -> class
-> IntMap.IntMap (Set.HashSet (EClassId, ENode)) -- ^ stored parent edges per class
-> IntMap.IntMap (Maybe Double, Maybe Double, Int, [Target]) -- ^ fit data
-> IntMap.IntMap Int -- ^ eid -> height
-> IntMap.IntMap EClassRow
buildClasses canon nodeToEClass storedParents fit heights =
IntMap.fromList [ (eid, mkRow eid) | (eid, c) <- IntMap.toList canon, c == eid ]
where
parentsOf :: IntMap.IntMap (Set.HashSet (EClassId, ENode))
parentsOf = IntMap.fromListWith Set.union
[ (c, Set.singleton (eid, en))
| (en, eid) <- HashMap.toList nodeToEClass
, c <- enodeChildren en ]
mkRow :: EClassId -> EClassRow
mkRow eid =
let nodes = Set.fromList [ en | (en, eid') <- HashMap.toList nodeToEClass, eid' == eid ]
h = IntMap.findWithDefault 0 eid heights
(fitM, dlM, sz, theta) = IntMap.findWithDefault (Nothing, Nothing, 0, []) eid fit
stored = IntMap.findWithDefault Set.empty eid storedParents
parents = if Set.null stored
then IntMap.findWithDefault Set.empty eid parentsOf
else stored
in EClassRow
{ _rcNodes = nodes
, _rcParents = parents
, _rcHeight = h
, _rcInfo = EData 0 (headOrDefault (EVar 0) (Set.toList nodes)) NotConst
fitM dlM theta sz }
headOrDefault :: a -> [a] -> a
headOrDefault def [] = def
headOrDefault _ (x:_) = x
-- | Push the graph's risk metrics into the @dataset_fit@ table for a dataset
-- in a legacy combined DB (leaves structure intact). On a paged graph the
-- classes are streamed page by page (O(1) memory) rather than materialized via
-- 'cpsAll'.
pushFit :: SqlBackend db => db -> Int -> EGraph -> IO ()
pushFit = pushFitWith createSchema
-- | Like 'pushFit', but creates the per-dataset fit schema. Use this when the
-- fit table lives in a separate fit DB and the graph is resident (or the page
-- store lives in the same DB).
pushFitFit :: SqlBackend db => db -> Int -> EGraph -> IO ()
pushFitFit = pushFitWith createSchemaFit
-- | Write fit rows to a separate fit DB, reading paged classes from the egraph
-- DB. Required when the graph is paged and the fit DB is a different file: the
-- paged page store lives in the egraph DB, not in the fit DB.
pushFitSplit :: (SqlBackend egDb, SqlBackend fitDb) => egDb -> fitDb -> Int -> EGraph -> IO ()
pushFitSplit = pushFitWithSource createSchemaFit
pushFitWith :: SqlBackend db => (db -> IO ()) -> db -> Int -> EGraph -> IO ()
pushFitWith mkSchema db = pushFitWithSource mkSchema db db
pushFitWithSource
:: (SqlBackend egDb, SqlBackend fitDb)
=> (fitDb -> IO ())
-> egDb -> fitDb -> Int -> EGraph -> IO ()
pushFitWithSource mkSchema egDb fitDb dsid eg = do
mkSchema fitDb
run fitDb "DELETE FROM dataset_fit WHERE dataset_id = ?" [SqlInteger (fromIntegral dsid)]
case _classStore eg of
Nothing -> do
let rows0 = exportEGraph eg
writeDatasetFitRows fitDb dsid rows0
Just _ -> do
-- flush so every dirty page is in @cstore_page@, then stream each class
-- from the egraph DB and write its row to the fit DB (O(1) memory).
flushStore eg
execDb fitDb "BEGIN"
streamPages egDb classStoreTable $ \_eid blob -> do
let ec = decode (BL.fromStrict blob)
info = _info ec
writeDatasetFit fitDb dsid (_eClassId ec) (_fitness info) (_dl info)
(T.pack (serializeTheta (_theta info))) (_size info)
execDb fitDb "COMMIT"
-- | Overwrite in-memory fitness/DL with the values currently stored in the
-- database for a dataset (per e-class, by canonical id).
refreshFitness :: SqlBackend db => db -> Int -> EGraph -> IO (Either String EGraph)
refreshFitness db dsid eg = do
dsFit <- readDatasetFit db dsid
let applyFit :: ClassStore m => EGraphST m ()
applyFit = forM_ dsFit $ \(eid, (fitM, _, _, theta)) ->
case fitM of
Nothing -> pure ()
Just f -> do
c <- canonical eid
insertFitness c f (parseTheta (T.unpack theta))
case _classStore eg of
-- Resident graph: the pure ClassStore instance is sufficient.
Nothing -> pure (Right (runIdentity $ execStateT applyFit eg))
-- Paged graph: run through the IO ClassStore so page reads/writes use the
-- live store (the pure instance would see an empty resident map).
Just _ -> do
(_, eg') <- runStateT applyFit eg
pure (Right eg')
-- | Bulk-load page blobs for a list of e-class IDs in one query, returning
-- an IntMap of deserialized EClass values. For lists exceeding SQLite's
-- parameter limit, the query is chunked automatically.
loadPagesBulk :: SqlBackend db => db -> [EClassId] -> IO (IntMap.IntMap EClass)
loadPagesBulk _ [] = pure IntMap.empty
loadPagesBulk db eids = do
let chunks = chunkList 500 eids
foldM (\acc chunkIds -> do
pages <- loadChunk chunkIds
pure $! IntMap.union acc pages) IntMap.empty chunks
where
chunkList _ [] = []
chunkList n xs = let (h, t) = splitAt n xs in h : chunkList n t
loadChunk ids = do
let placeholders = intercalate "," (replicate (length ids) "?")
params = map (SqlInteger . fromIntegral) ids
rows <- queryDb db
("SELECT key, blob FROM cstore_page WHERE key IN (" <> T.pack placeholders <> ")")
params
pure $ IntMap.fromList
[ (sqlToInt k, decode (BL.fromStrict (sqlToBlob b)))
| [k, b] <- rows
]