diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,11 @@ +# Changelog for srtree-db + +## 0.1.0.0 + +- Initial release +- Two-section data model: dataset-agnostic e-graph section + per-dataset fit section +- SQLite and PostgreSQL backends +- Out-of-core import with page store +- Dataset-aware queries: topN, pareto, paretoBySize, distributionCounts +- Frontier re-saturation support +- Streaming page export via cursor (O(1) memory) diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +Copyright (c) 2026, folivetti + + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * 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. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"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 +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 +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,414 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE BangPatterns #-} + +-- | Out-of-core feasibility spike. +-- +-- Measures the cost of paging individual e-classes in and out of a SQL +-- database under the two access shapes the e-graph algorithm produces, and +-- compares against a RAM-resident map. +-- +-- * *worklist* - 90% of accesses hit a 10% "hot" set (congruence/rebuild +-- locality), 10% are cold. +-- * *random* - uniform random classes (worst case, cache-hostile). +-- +-- Each class is stored as one row (a binary page). A tiny FIFO cache models +-- the future ClassStore; dirty pages are coalesced and flushed in batches. +-- +-- Usage: Bench [nClasses] [ops] (defaults 100000 200000) +-- Env: PGDSN= (optional; adds PostgreSQL cells) + +module Main (main) where + +import Control.Exception (bracket, catch, SomeException) +import Control.Monad (forM, forM_, when) +import Data.IORef +import qualified Data.IntMap.Strict as IM +import qualified Data.Map.Strict as Map +import qualified Data.ByteString as BS +import qualified Data.Text as T +import Data.Time.Clock (getCurrentTime, diffUTCTime) +import System.Environment (lookupEnv, getArgs) +import System.Directory (removeFile) + +import Database.SQLite3 (Database, open, close) +import Database.PostgreSQL.LibPQ (Connection) + +import Algorithm.EqSat.Storage.Backend (SqlValue(..), SqlBackend(..)) +import Algorithm.EqSat.Storage.SQLite () +import Algorithm.EqSat.Storage.Postgres (connectPostgres, closePostgres) + +-- --------------------------------------------------------------------------- +-- pseudo-random (LCG); deterministic and dependency-free + +next :: Int -> Int +next r = (r * 1103515245 + 12345) `mod` 2147483647 + +pageBytes :: Int -> BS.ByteString +pageBytes eid = BS.pack (take 512 (byteLoop (eid * 2654435761 + 7))) + where + byteLoop r = let r' = next r in fromIntegral (r' `mod` 256) : byteLoop r' + +pageHex :: Int -> T.Text +pageHex eid = T.pack (concatMap go (BS.unpack (pageBytes eid))) + where + go b = let hi = fromIntegral (b `div` 16); lo = fromIntegral (b `mod` 16) + in "0123456789abcdef" !! hi : ["0123456789abcdef" !! lo] + +-- --------------------------------------------------------------------------- +-- LRU page cache +-- +-- * @cMap@ - eid -> (content, last-access tick) +-- * @cClock@ - tick -> eid (ordered index for O(log n) eviction) +-- * @cPend@ - dirty pages awaiting write-back +-- All operations are O(log n). + +data Cache = Cache + { cMap :: !(IM.IntMap (BS.ByteString, Int)) -- eid -> (content, tick) + , cClock :: !(Map.Map Int Int) -- tick -> eid + , cPend :: !(IM.IntMap BS.ByteString) -- dirty pages awaiting write-back + , cSize :: !Int -- |cMap| (O(1), avoid IM.size) + , cPendN :: !Int -- |cPend| (O(1), avoid IM.size) + , cTick :: !Int + } + +emptyCache :: Cache +emptyCache = Cache IM.empty Map.empty IM.empty 0 0 0 + +-- | Record an access to @eid@ (content known, page resident). +touch :: Int -> BS.ByteString -> Cache -> Cache +touch eid page cc = + let t = cTick cc + 1 + wasResident = IM.member eid (cMap cc) + old = snd <$> IM.lookup eid (cMap cc) + rest = IM.insert eid (page, t) (cMap cc) + clk' = maybe (cClock cc) (`Map.delete` cClock cc) old + sz' = if wasResident then cSize cc else cSize cc + 1 + in cc { cMap = rest, cClock = Map.insert t eid clk', cTick = t, cSize = sz' } + +-- | Insert (or re-touch) @eid@, evicting the LRU page when at capacity. +-- Evicted dirty pages stay in @cPend@ until flushed. +insertEvict :: Int -> Int -> BS.ByteString -> Cache -> Cache +insertEvict cap eid page cc + | cap <= 0 = cc + | cSize cc < cap = touch eid page cc + | otherwise = case Map.lookupMin (cClock cc) of + Nothing -> touch eid page cc + Just (tOld, ev) -> + let cc' = cc { cMap = IM.delete ev (cMap cc) + , cClock = Map.delete tOld (cClock cc) + , cSize = cSize cc - 1 } + in touch eid page cc' + +markDirty :: Int -> Cache -> Cache +markDirty eid cc = + case IM.lookup eid (cMap cc) of + Nothing -> cc + Just (bs, _) -> + let wasDirty = IM.member eid (cPend cc) + in cc { cPend = IM.insert eid bs (cPend cc) + , cPendN = if wasDirty then cPendN cc else cPendN cc + 1 } + +-- --------------------------------------------------------------------------- +-- workload runners + +type LoadFn = Int -> IO BS.ByteString +type FlushFn = [(Int, BS.ByteString)] -> IO () + +-- | Run @nOps@ accesses under the given access shape; returns hit count. +runPaged :: Int -> Int -> Int -> Int -> Double -> Double -> Bool + -> LoadFn -> FlushFn -> IO Int +runPaged cap nOps n flushEvery hotProb rmwProb useWorklist load flush = do + cacheRef <- newIORef emptyCache + hitsRef <- newIORef (0 :: Int) + let nHot = max 1 (n `div` 10) + flushNow = do + c <- readIORef cacheRef + let pend = IM.toList (cPend c) + when (not (null pend)) $ do + flush pend + modifyIORef' cacheRef (\cc -> cc { cPend = IM.empty, cPendN = 0 }) + step r = do + let r1 = next r + r2 = next r1 + r3 = next r2 + p1 = fromIntegral (r2 `mod` 1000) / 1000 + p2 = fromIntegral (r3 `mod` 1000) / 1000 + eid = if useWorklist && p1 < hotProb + then r1 `mod` nHot + else r3 `mod` n + rmw = p2 < rmwProb + c0 <- readIORef cacheRef + if IM.member eid (cMap c0) + then do + modifyIORef' hitsRef (+1) + when rmw (modifyIORef' cacheRef (markDirty eid)) + pure r3 + else do + page <- load eid + modifyIORef' cacheRef (insertEvict cap eid page) + when rmw (modifyIORef' cacheRef (markDirty eid)) + c1 <- readIORef cacheRef + when (cPendN c1 >= flushEvery) flushNow + pure r3 + go !acc r + | acc <= 0 = pure () + | otherwise = step r >>= go (acc - 1) + _ <- go nOps 12345 + flushNow + readIORef hitsRef + +-- | RAM baseline: all pages resident; every op is an IntMap lookup/update. +runRAM :: Int -> Int -> Double -> IO () +runRAM n nOps rmwProb = do + ref <- newIORef (IM.fromList [ (i, pageBytes i) | i <- [0 .. n - 1] ]) + let step r = do + let r1 = next r + eid = r1 `mod` n + modifyIORef' ref (IM.alter (fmap modPage) eid) + pure (next r1) + modPage b = BS.take 1 b <> BS.drop 1 b + go !acc r + | acc <= 0 = pure () + | otherwise = step r >>= go (acc - 1) + _ <- go nOps 12345 + pure () + +-- --------------------------------------------------------------------------- +-- table setup / bulk load + +sqliteDDL, pgDDL :: T.Text +sqliteDDL = "CREATE TABLE IF NOT EXISTS kpage (eid INTEGER PRIMARY KEY, blob TEXT NOT NULL)" +pgDDL = "CREATE TABLE IF NOT EXISTS kpage (eid BIGINT PRIMARY KEY, blob TEXT NOT NULL)" + +setupSQLite :: Database -> IO () +setupSQLite db = do + execDb db "DROP TABLE IF EXISTS kpage" + execDb db "PRAGMA journal_mode=WAL" + execDb db sqliteDDL + +setupPG :: Connection -> IO () +setupPG conn = do + execDb conn "DROP TABLE IF EXISTS kpage" + execDb conn pgDDL + +bulkLoad :: SqlBackend db => db -> Int -> IO () +bulkLoad db n = do + execDb db "BEGIN" + let chunk lo = forM_ [lo .. lo + 4999] $ \i -> + runDb db "INSERT INTO kpage (eid, blob) VALUES (?, ?)" + [ SqlInteger (fromIntegral i), SqlText (pageHex i) ] + forM_ [0, 5000 .. n - 1] chunk + execDb db "COMMIT" + +mkLoad :: SqlBackend db => db -> Int -> IO BS.ByteString +mkLoad db eid = do + rows <- queryDb db "SELECT blob FROM kpage WHERE eid = ?" [SqlInteger (fromIntegral eid)] + case rows of + [[SqlText t]] -> pure (unhex (T.unpack t)) + _ -> fail ("missing page " <> show eid) + +unhex :: String -> BS.ByteString +unhex = BS.pack . go + where + go (a:b:r) = fromIntegral (hex a * 16 + hex b) : go r + go _ = [] + hex c | c >= '0' && c <= '9' = fromEnum c - fromEnum '0' + | otherwise = fromEnum c - fromEnum 'a' + 10 + +-- | Write back a batch of dirty pages in one transaction. +mkFlush :: SqlBackend db => db -> [(Int, BS.ByteString)] -> IO () +mkFlush db pend = do + execDb db "BEGIN" + forM_ pend $ \(eid, _bs) -> + runDb db "UPDATE kpage SET blob = blob WHERE eid = ?" [SqlInteger (fromIntegral eid)] + execDb db "COMMIT" + +-- --------------------------------------------------------------------------- +-- driver + +data Cell = Cell String Double Double -- label, ops/sec, hit% + +timeIO :: IO a -> IO (a, Double) +timeIO act = do + t0 <- getCurrentTime + x <- act + t1 <- getCurrentTime + pure (x, realToFrac (diffUTCTime t1 t0)) + +runRAMCase :: String -> Int -> Int -> Double -> IO Cell +runRAMCase lbl n ops rmw = do + (_, s) <- timeIO (runRAM n ops rmw) + pure (Cell lbl (fromIntegral ops / s) 100) + +runPgCase :: String -> Int -> Int -> Double -> Bool -> IO Cell +runPgCase lbl n ops rmw worklist = do + dsn <- lookupEnv "PGDSN" + case dsn of + Nothing -> pure (Cell (lbl <> " [skipped]") 0 0) + Just d -> bracket (connectPostgres d) closePostgres $ \conn -> do + setupPG conn + bulkLoad conn n + (hits, s) <- timeIO $ runPaged (n `div` 4) ops n 2000 0.9 rmw worklist + (mkLoad conn) (mkFlush conn) + pure (Cell lbl (fromIntegral ops / s) (100 * fromIntegral hits / fromIntegral ops)) + +runSQLiteCase :: String -> Int -> Int -> Double -> Int -> Bool -> IO Cell +runSQLiteCase lbl n ops rmw cap worklist = do + let path = "/tmp/opencode/bench.sqlite" + removeFile path `catch` (\(_ :: SomeException) -> pure ()) + bracket (open (T.pack path)) close $ \db -> do + setupSQLite db + bulkLoad db n + (hits, s) <- timeIO $ runPaged cap ops n 2000 0.9 rmw worklist + (mkLoad db) (mkFlush db) + pure (Cell lbl (fromIntegral ops / s) (100 * fromIntegral hits / fromIntegral ops)) + +-- | Same as the SQLite case but with a trivial in-memory "load": isolates +-- pure cache overhead from database I/O. +runNullCase :: String -> Int -> Int -> Double -> Int -> Bool -> IO Cell +runNullCase lbl n ops rmw cap worklist = do + let page0 = pageBytes 0 -- pre-forced; cache never rebuilds it + (hits, s) <- timeIO $ runPaged cap ops n 2000 0.9 rmw worklist + (\_ -> pure page0) (\_ -> pure ()) + pure (Cell lbl (fromIntegral ops / s) (100 * fromIntegral hits / fromIntegral ops)) + +-- | Microbenchmarks of the container primitives used by the LRU cache. +microMaps :: Int -> IO Cell +microMaps n = do + let bigIM = IM.fromList [ (i, (BS.replicate 512 0, i)) | i <- [0 .. n - 1] ] + bigM = Map.fromList [ (i, i) | i <- [0 .. n - 1] ] + loopIM !acc !i m = if acc <= 0 then pure () else loopIM (acc - 1) (i + 1) (IM.insert i (BS.replicate 512 0, i) m) + loopM !acc !i m = if acc <= 0 then pure () else loopM (acc - 1) (i + 1) (Map.insert i i m) + loopMem !acc !i m = if acc <= 0 then pure () else loopMem (acc - 1) (i + 1) (if IM.member (i `mod` n) m then m else m) + (_, s1) <- timeIO (loopIM 2000000 0 bigIM) + (_, s2) <- timeIO (loopM 2000000 0 bigM) + (_, s3) <- timeIO (loopMem 5000000 0 bigIM) + -- replicate 'touch' on a ~full cache (insert + clock insert + eviction) + let touchLoop !acc i (m :: IM.IntMap (BS.ByteString, Int)) (clk :: Map.Map Int Int) = + if acc <= 0 then pure () else + let t = i + eid = i `mod` n + (m2, clk2) = case Map.lookupMin clk of + Nothing -> (IM.insert eid (BS.replicate 512 0, t) m, Map.insert t eid clk) + Just (tOld, ev) -> + ( IM.insert eid (BS.replicate 512 0, t) (IM.delete ev m) + , Map.insert t eid (Map.delete tOld clk) ) + in touchLoop (acc - 1) (i + 1) m2 clk2 + (_, s4) <- timeIO (touchLoop 2000000 0 bigIM (Map.fromList [(i, i) | i <- [0 .. 49999]])) + -- exact replication of runPaged's miss branch on an IORef'd Cache + let cap = 25000 + mkPage = BS.replicate 512 0 + missLoop !acc i = if acc <= 0 then pure () else do + ref <- newIORef emptyCache + let one ref i = do + c0 <- readIORef ref + let eid = i `mod` 100000 + hit = IM.member eid (cMap c0) + if hit + then modifyIORef' ref (markDirty eid) + else do + modifyIORef' ref (insertEvict cap eid mkPage) + c1 <- readIORef ref + when (cPendN c1 >= 2000) (pure ()) + pure (i + 1) + one ref i >>= missLoop (acc - 1) + (_, s5) <- timeIO (missLoop 200000 0) + -- same miss branch but with a persistent IORef'd cache (grows to full) + let missLoopP !acc ref = if acc <= 0 then pure () else do + c0 <- readIORef ref + let eid = (acc + 1000) `mod` 100000 + hit = IM.member eid (cMap c0) + if hit + then modifyIORef' ref (markDirty eid) + else do + modifyIORef' ref (insertEvict cap eid mkPage) + c1 <- readIORef ref + when (cPendN c1 >= 2000) (pure ()) + missLoopP (acc - 1) ref + pRef <- newIORef emptyCache + (_, s5p) <- timeIO (missLoopP 200000 pRef) + -- verbatim copy of runPaged's step/go, cache enabled vs disabled + let mkPage = BS.replicate 512 0 + stepW cap' cacheRef hitsRef r = do + let r1 = next r; r2 = next r1; r3 = next r2 + p1 = fromIntegral (r2 `mod` 1000) / 1000 + p2 = fromIntegral (r3 `mod` 1000) / 1000 + eid = r3 `mod` 100000 + rmw = p2 < 0.3 + c0 <- readIORef cacheRef + if IM.member eid (cMap c0) + then do + modifyIORef' hitsRef (+1) + when rmw (modifyIORef' cacheRef (markDirty eid)) + pure r3 + else do + let page = mkPage + modifyIORef' cacheRef (insertEvict cap' eid page) + when rmw (modifyIORef' cacheRef (markDirty eid)) + c1 <- readIORef cacheRef + when (cPendN c1 >= 2000) (pure ()) + pure r3 + goW cacheRef hitsRef !acc r = if acc <= 0 then pure () else stepW 25000 cacheRef hitsRef r >>= goW cacheRef hitsRef (acc - 1) + goW0 cacheRef hitsRef !acc r = if acc <= 0 then pure () else stepW 0 cacheRef hitsRef r >>= goW0 cacheRef hitsRef (acc - 1) + runW goF = do + cacheRef <- newIORef emptyCache + hitsRef <- newIORef (0 :: Int) + goF cacheRef hitsRef 200000 12345 + (_, s6) <- timeIO (runW goW) + (_, s7) <- timeIO (runW goW0) + pure (Cell ("micro-IM.insert(" <> show (round (2000000 / s1)) <> "ops/s) Map.insert(" + <> show (round (2000000 / s2)) <> ") IM.member(" + <> show (round (5000000 / s3)) <> ") touch(" + <> show (round (2000000 / s4)) <> ") missbranch(" + <> show (round (200000 / s5)) <> ") missbranchP(" + <> show (round (200000 / s5p)) <> ") step-cap25(" + <> show (round (200000 / s6)) <> ") step-cap0(" + <> show (round (200000 / s7)) <> ")") 0 0) + +-- | Bare loop floor: 5M iterations of LCG + IORef counter, nothing else. +selfbench :: IO Cell +selfbench = do + ref <- newIORef (0 :: Int) + let go !acc !r + | acc <= 0 = pure () + | otherwise = go (acc - 1) (next r) + go2 !acc !r + | acc <= 0 = pure () + | otherwise = do + let r1 = next r + modifyIORef' ref (+1) + go2 (acc - 1) r1 + (_, s1) <- timeIO (go 5000000 12345) + (_, s2) <- timeIO (go2 5000000 12345) + pure (Cell ("selfbench-loop(" <> show (round (5000000 / s1)) <> ",withIORef=" <> show (round (5000000 / s2)) <> ")") 0 0) + +main :: IO () +main = do + args <- getArgs + let n = case args of (a:_) -> read a; _ -> 100000 + ops = case args of (_:b:_) -> read b; _ -> 200000 + rmw = 0.3 + q1 = n `div` 4 + q2 = n `div` 2 + putStrLn ("# classes=" <> show n <> " ops=" <> show ops <> " flushEvery=2000 rmw=0.3") + putStrLn "# cell\tops/sec\thit%" + sb <- selfbench + putStrLn (let Cell l _ _ = sb in l) + mm <- microMaps 100000 + putStrLn (let Cell l _ _ = mm in l) + cRam <- runRAMCase "ram-random-resident" n ops rmw + cSu <- runSQLiteCase "sqlite-uncached-random" n ops rmw 0 False + cSw10 <- runSQLiteCase "sqlite-worklist-cap25" n ops rmw q1 True + cSw50 <- runSQLiteCase "sqlite-worklist-cap50" n ops rmw q2 True + cSr <- runSQLiteCase "sqlite-random-cap25" n ops rmw q1 False + cR0 <- runSQLiteCase "sqlite-worklist-cap25-normw" n ops 0 q1 True + cR5 <- runSQLiteCase "sqlite-worklist-cap25-rmw50" n ops 0.5 q1 True + cN0 <- runNullCase "null-worklist-cap25" n ops rmw q1 True + cN1 <- runNullCase "null-random-cap25" n ops rmw q1 False + cN2 <- runNullCase "null-random-cap0" n ops rmw 0 False + cN3 <- runNullCase "null-random-cap100" n ops rmw n False + cPu <- runPgCase "pg-uncached-random" n ops rmw False + cPw <- runPgCase "pg-worklist-cap25" n ops rmw True + mapM_ (\(Cell l v h) -> putStrLn (l <> "\t" <> show (round v) <> "\t" <> show (round h))) + [cRam, cSu, cSw10, cSw50, cSr, cR0, cR5, cN0, cN1, cN2, cN3, cPu, cPw] diff --git a/src/Algorithm/EqSat/Storage/Backend.hs b/src/Algorithm/EqSat/Storage/Backend.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Backend.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Driver-neutral SQL backend for srtree e-graph persistence. +-- +-- The storage layer ('Algorithm.EqSat.Storage.SQLite' and +-- 'Algorithm.EqSat.Storage.Postgres') is written against this tiny interface +-- instead of a concrete driver, so the same serialization, import and query +-- code runs on SQLite and PostgreSQL. Drivers only provide: +-- +-- * 'execDb' - statements without parameters (DDL, BEGIN/COMMIT, DELETE) +-- * 'runDb' - parameterized statements that return no rows (INSERT/UPDATE) +-- * 'queryDb' - parameterized statements returning rows (SELECT) +-- * 'createSchemaDb' - create the driver-specific schema +-- +-- Parameters use the positional @?@ placeholder in the shared SQL; the +-- PostgreSQL driver rewrites them to @$n@. NULL is expressed as the literal +-- @NULL@ in the shared SQL (never as a parameter), so drivers do not need a +-- NULL parameter representation. +module Algorithm.EqSat.Storage.Backend + ( SqlValue(..) + , SqlBackend(..) + , sqlToInt + , sqlToMaybeDouble + , sqlToText + ) where + +import Data.Int (Int64) +import Data.Maybe (fromMaybe, listToMaybe) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.ByteString as BS + +import Algorithm.EqSat.Egraph (EClassId) + +-- | A driver-neutral value bound to a @?@ parameter or returned by a query. +data SqlValue = SqlInteger Int64 + | SqlFloat Double + | SqlText Text + | SqlNull + deriving (Eq, Show) + +-- | The minimal SQL surface used by the storage layer. +class SqlBackend db where + -- | Execute a statement without parameters. + execDb :: db -> Text -> IO () + -- | Execute a parameterized statement that returns no rows. + runDb :: db -> Text -> [SqlValue] -> IO () + -- | Execute an idempotent insert, ignoring any row that would violate a + -- primary/unique constraint (re-seeing a node already present). The @Text@ + -- argument is the @table (cols) VALUES (?,...)@ tail *without* the leading + -- @INSERT INTO@ and without a conflict clause; each driver supplies its own + -- native prefix/suffix (SQLite @INSERT OR IGNORE INTO@, Postgres @ON + -- CONFLICT DO NOTHING@). + insertIgnore :: db -> Text -> [SqlValue] -> IO () + -- | Run a parameterized query and return the raw result grid. + queryDb :: db -> Text -> [SqlValue] -> IO [[SqlValue]] + -- | Stream (bounded) the distinct e-class ids whose e-class contains a node + -- with the given @op_detail@, for 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). Drivers that expose a cursor + -- ('Database.SQLite3') implement this O(1)-memory; others fall back to a grid + -- 'queryDb' (unbounded, documented). + streamByOp :: db -> Text -> Int -> [EClassId] -> IO [EClassId] + -- | Stream the @key, blob@ rows of a key-value page table (e.g. + -- @cstore_page@) to a callback, one at a time, so a full pass over every page + -- (e.g. 'pushFit') stays O(1) memory. Drivers with a cursor ('Database.SQLite3') + -- implement this; others fall back to a grid 'queryDb' (unbounded, documented). + -- The blob is delivered hex-decoded (raw) to the callback. + streamPages :: db -> Text -> (Int64 -> BS.ByteString -> IO ()) -> IO () + -- | Create the schema (tables, indexes) for this driver. + createSchemaDb :: db -> IO () + +sqlToInt :: SqlValue -> Int +sqlToInt (SqlInteger n) = fromIntegral n +sqlToInt (SqlText t) = fromMaybe 0 (listToMaybe [ i | (i, "") <- reads (T.unpack t) ]) +sqlToInt _ = 0 + +sqlToMaybeDouble :: SqlValue -> Maybe Double +sqlToMaybeDouble SqlNull = Nothing +sqlToMaybeDouble (SqlFloat d) = Just d +sqlToMaybeDouble (SqlInteger n) = Just (fromIntegral n) +sqlToMaybeDouble (SqlText t) = case reads (T.unpack t) of + [(d, "")] -> Just d + _ -> Nothing +sqlToMaybeDouble _ = Nothing + +sqlToText :: SqlValue -> Text +sqlToText (SqlText t) = t +sqlToText (SqlInteger n) = T.pack (show n) +sqlToText (SqlFloat d) = T.pack (show d) +sqlToText SqlNull = "" diff --git a/src/Algorithm/EqSat/Storage/ClassStore.hs b/src/Algorithm/EqSat/Storage/ClassStore.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/ClassStore.hs @@ -0,0 +1,441 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE BangPatterns #-} + +-- | A lazily loaded, LRU-cached, write-back page store for e-class blobs. +-- +-- Provides the storage substrate for an out-of-core e-graph: individual +-- e-classes (serialized to 'ByteString' pages) live in a single key-value +-- table and are paged in and out of a bounded LRU cache. Dirty pages are +-- coalesced and flushed in batches, so the write-back cost is a small number +-- of transactions instead of one per mutation. +-- +-- Design (validated by the 'srtree-db-bench' spike): +-- +-- * LRU cache with an O(log n) ordered recency index; cache capacity is an +-- explicit bound (a fraction of the total class count). +-- * O(1) resident/pending counters (never 'Data.IntMap.size', which is +-- O(n) per call). +-- * Every access refreshes recency (true LRU), so the hot classes a +-- rebuild/rewrite pass revisits stay resident. +-- * Pages are stored hex-encoded in a TEXT column so the same DDL/CRUD +-- runs unchanged on SQLite and PostgreSQL (a bytea/BLOB column is a +-- possible later optimization). +-- +-- Driver-neutrality note: the store talks only through 'SqlBackend' and +-- spells writes as DELETE+INSERT inside one transaction (both drivers +-- support these statements; no driver-specific upsert syntax). +-- +-- The table is a plain key-value page store; the actual +-- 'Algorithm.EqSat.Egraph.EClass' serialization and the parent relation are +-- layered on top of it by the caller. + +module Algorithm.EqSat.Storage.ClassStore + ( PageStore + , newPageStore + , readPage + , writePage + , deletePage + , writeback + , pendingCount + , residentCount + , flushEvery + , allPages + , classStoreTable + , openClassStore + , classStoreHandle + , frontierTable + , markFrontier + , loadFrontierRows + , clearFrontier + , setFrontierActive + , initFrontier + , hex + , unhex + ) where + +import Control.Monad (forM_, unless, when) +import Data.IORef +import qualified Data.IntMap.Strict as IM +import qualified Data.IntSet as IntSet +import qualified Data.Map.Strict as Map +import qualified Data.ByteString as BS +import qualified Data.HashSet as Set +import qualified Data.Text as T +import Data.Binary (encode, decode) +import Data.List (nub) +import qualified Data.ByteString.Lazy as BL + +import Algorithm.EqSat.Storage.Backend + ( SqlValue(..), SqlBackend(..), sqlToInt, sqlToText ) +import Algorithm.EqSat.Storage.Types + ( enodeKey, enodeOpTag, enodeOpDetail, opDetailOf ) +import Algorithm.EqSat.Egraph + ( EClass, EClassPageStore(..), ENode(..), _eClassId ) + +-- --------------------------------------------------------------------------- +-- hex encoding of page blobs (driver-neutral TEXT storage) + +hex :: BS.ByteString -> T.Text +hex = T.pack . concatMap go . BS.unpack + where + go b = + let hi = fromIntegral (b `div` 16) + lo = fromIntegral (b `mod` 16) + in "0123456789abcdef" !! hi : ["0123456789abcdef" !! lo] + +unhex :: T.Text -> BS.ByteString +unhex = BS.pack . go . T.unpack + where + go (a:b:r) = fromIntegral (hexv a * 16 + hexv b) : go r + go _ = [] + hexv c | c >= '0' && c <= '9' = fromEnum c - fromEnum '0' + | otherwise = fromEnum c - fromEnum 'a' + 10 + +-- --------------------------------------------------------------------------- +-- LRU page cache + +data PageCache = PageCache + { pcMap :: !(IM.IntMap (BS.ByteString, Int)) -- eid -> (page, recency tick) + , pcClock :: !(Map.Map Int Int) -- tick -> eid (for eviction) + , pcPend :: !(IM.IntMap BS.ByteString) -- dirty pages awaiting write-back + , pcSize :: !Int -- |pcMap| (O(1)) + , pcPendN :: !Int -- |pcPend| (O(1)) + , pcTick :: !Int + , pcCap :: !Int + } + +emptyCache :: Int -> PageCache +emptyCache cap = PageCache IM.empty Map.empty IM.empty 0 0 0 cap + +-- | Record an access to @eid@ with the given page content, updating recency. +touch :: Int -> BS.ByteString -> PageCache -> PageCache +touch eid page pc = + let t = pcTick pc + 1 + resident = IM.member eid (pcMap pc) + old = snd <$> IM.lookup eid (pcMap pc) + m' = IM.insert eid (page, t) (pcMap pc) + clk' = maybe (pcClock pc) (`Map.delete` pcClock pc) old + sz' = if resident then pcSize pc else pcSize pc + 1 + in pc { pcMap = m', pcClock = Map.insert t eid clk', pcSize = sz', pcTick = t } + +-- | Insert (or re-touch) @eid@, evicting the least-recently-used page when at +-- capacity. Evicted dirty pages stay in @pcPend@ until written back. +-- A capacity of 0 keeps every page resident (unbounded). +insertEvict :: Int -> BS.ByteString -> PageCache -> PageCache +insertEvict eid page pc + | pcCap pc <= 0 = touch eid page pc + | pcSize pc < pcCap pc = touch eid page pc + | otherwise = case Map.lookupMin (pcClock pc) of + Nothing -> touch eid page pc + Just (tOld, ev) -> + touch eid page + pc { pcMap = IM.delete ev (pcMap pc) + , pcClock = Map.delete tOld (pcClock pc) + , pcSize = pcSize pc - 1 } + +-- | Mark @eid@ dirty (ensuring it is resident first). +markDirty :: Int -> BS.ByteString -> PageCache -> PageCache +markDirty eid page pc = + let pc' = if IM.member eid (pcMap pc) then pc else insertEvict eid page pc + wasDirty = IM.member eid (pcPend pc') + in pc' { pcPend = IM.insert eid page (pcPend pc') + , pcPendN = if wasDirty then pcPendN pc' else pcPendN pc' + 1 } + +-- --------------------------------------------------------------------------- +-- store + +-- | A paging store over a table @key TEXT PRIMARY KEY, blob TEXT NOT NULL@, +-- connected to a 'SqlBackend' database. +data PageStore db = PageStore + { psDb :: db + , psTable :: !T.Text + , psCache :: !(IORef PageCache) + , psNodes :: !(IORef (Set.HashSet (Int, ENode))) -- newly-created nodes awaiting write-back + , psCanon :: !(IORef (IM.IntMap Int)) -- pending canonical (eid -> representative) rows + , psFrontier :: !(IORef IntSet.IntSet) -- recently-changed e-classes (dirty set) + , psFrontierActive :: !(IORef Bool) -- restrict the matcher to the frontier + , psCap :: !Int + , psFlushEvery :: !Int + } + +-- | Open a page store: creates the (driver-neutral) page table if missing, +-- with the given bounded cache (in pages) and a write-back flush every +-- @flushEvery@ dirty pages. +-- +-- The table name is a plain identifier (never interpolate user input). +newPageStore :: SqlBackend db => db -> T.Text -> Int -> Int -> IO (PageStore db) +newPageStore db tbl cap flushEvery' = do + execDb db + ("CREATE TABLE IF NOT EXISTS " <> tbl <> + " (key TEXT PRIMARY KEY, blob TEXT NOT NULL)") + cache <- newIORef (emptyCache cap) + nodes <- newIORef Set.empty + canons <- newIORef IM.empty + frontier <- newIORef IntSet.empty + active <- newIORef False + pure (PageStore db tbl cache nodes canons frontier active cap flushEvery') + +-- | Read the page for @eid@ (cache miss loads from the database). Refreshes +-- LRU recency on hit so hot classes stay resident. +readPage :: SqlBackend db => PageStore db -> Int -> IO (Maybe BS.ByteString) +readPage ps eid = do + c0 <- readIORef (psCache ps) + case IM.lookup eid (pcMap c0) of + Just (page, _) -> do + writeIORef (psCache ps) (touch eid page c0) + pure (Just page) + Nothing -> case IM.lookup eid (pcPend c0) of + -- the page is dirty (written but not yet flushed) and was evicted from + -- the LRU cache; it is NOT in the database yet, so return it from the + -- pending set and restore it as resident. + Just page -> do + writeIORef (psCache ps) (insertEvict eid page c0) + pure (Just page) + Nothing -> do + rows <- queryDb (psDb ps) + ("SELECT blob FROM " <> psTable ps <> " WHERE key = ?") + [SqlInteger (fromIntegral eid)] + case rows of + [] -> pure Nothing + [[SqlText hv]] -> do + let page = unhex hv + writeIORef (psCache ps) (insertEvict eid page c0) + pure (Just page) + _ -> fail "ClassStore.readPage: unexpected row shape" + +-- | Write the page for @eid@ (resident + marked dirty). Triggers a batched +-- write-back once @flushEvery@ dirty pages have accumulated on a write. +writePage :: SqlBackend db => PageStore db -> Int -> BS.ByteString -> IO () +writePage ps eid page = do + modifyIORef' (psCache ps) (markDirty eid page) + c <- readIORef (psCache ps) + when (pcPendN c >= psFlushEvery ps) (writeback ps >> pure ()) + +-- | Remove @eid@ from the cache and the database. +deletePage :: SqlBackend db => PageStore db -> Int -> IO () +deletePage ps eid = do + modifyIORef' (psCache ps) $ \pc -> + pc { pcMap = IM.delete eid (pcMap pc) + , pcPend = IM.delete eid (pcPend pc) + , pcPendN = if IM.member eid (pcPend pc) then pcPendN pc - 1 else pcPendN pc } + runDb (psDb ps) + ("DELETE FROM " <> psTable ps <> " WHERE key = ?") + [SqlInteger (fromIntegral eid)] + +-- | Write back all pending dirty pages in one transaction. Returns the number +-- of pages written. +writeback :: SqlBackend db => PageStore db -> IO Int +writeback ps = do + c <- readIORef (psCache ps) + let pend = IM.toList (pcPend c) + unless (null pend) $ do + let db = psDb ps + tbl = psTable ps + execDb db "BEGIN" + forM_ pend $ \(eid, page) -> do + runDb db ("DELETE FROM " <> tbl <> " WHERE key = ?") + [SqlInteger (fromIntegral eid)] + runDb db ("INSERT INTO " <> tbl <> " (key, blob) VALUES (?, ?)") + [ SqlInteger (fromIntegral eid), SqlText (hex page) ] + execDb db "COMMIT" + modifyIORef' (psCache ps) (\cc -> cc { pcPend = IM.empty, pcPendN = 0 }) + flushNodes ps + flushCanon ps + pure (length pend) + +-- | Flush any newly-created nodes recorded via 'cpsRecordNode' into the +-- relational @enode@/@eclass_node@/@enode_child@ tables (idempotently), so the +-- streaming matcher's @op_detail@ index and the @enode_child@ children table +-- reflect the live graph. Runs inside its own transaction (the page write-back +-- above may be a no-op). +flushNodes :: SqlBackend db => PageStore db -> IO () +flushNodes ps = do + nodes <- readIORef (psNodes ps) + unless (Set.null nodes) $ do + let db = psDb ps + execDb db "BEGIN" + forM_ (Set.toList nodes) $ \(eid, en) -> do + let key = T.pack (enodeKey en) + insertIgnore db "enode (key, op, op_detail) VALUES (?, ?, ?)" + [ SqlText key + , SqlText (T.pack (enodeOpTag en)) + , SqlText (T.pack (enodeOpDetail en)) ] + insertIgnore db "eclass_node (eid, enode_key) VALUES (?, ?)" + [ SqlInteger (fromIntegral eid), SqlText key ] + -- ENAry children are stored relationally; keep them populated during eqsat + -- too (they were previously written only at import/save). + forM_ (naryChildren en) $ \(c, n) -> + insertIgnore db "enode_child (enode_key, child_eid, cnt) VALUES (?, ?, ?)" + [ SqlText key, SqlInteger (fromIntegral c), SqlInteger (fromIntegral n) ] + execDb db "COMMIT" + modifyIORef' (psNodes ps) (const Set.empty) + +-- | Children of an ENAry node as (class, multiplicity); empty otherwise. +naryChildren :: ENode -> [(Int, Int)] +naryChildren (ENAry _ m) = IM.toList m +naryChildren _ = [] + +-- | Flush any pending canonical rows (recorded via 'cpsRecordCanonical') into +-- @eclass.canonical@ (idempotent upsert), so the streaming canonical lookup +-- ('cpsCanonicalOf') reflects merges and new classes. Runs in its own +-- transaction. +flushCanon :: SqlBackend db => PageStore db -> IO () +flushCanon ps = do + canons <- readIORef (psCanon ps) + unless (IM.null canons) $ do + let db = psDb ps + execDb db "BEGIN" + forM_ (IM.toList canons) $ \(eid, canon) -> + runDb db ("INSERT INTO eclass (eid, canonical, height) VALUES (?, ?, 0) " + <> "ON CONFLICT(eid) DO UPDATE SET canonical=excluded.canonical") + [ SqlInteger (fromIntegral eid) + , SqlInteger (fromIntegral canon) ] + execDb db "COMMIT" + writeIORef (psCanon ps) IM.empty + +-- | Number of dirty pages not yet written back. +pendingCount :: PageStore db -> IO Int +pendingCount ps = pcPendN <$> readIORef (psCache ps) + +-- | Number of resident (cached) pages. +residentCount :: PageStore db -> IO Int +residentCount ps = pcSize <$> readIORef (psCache ps) + +-- | The configured flush threshold. +flushEvery :: PageStore db -> Int +flushEvery = psFlushEvery + +-- --------------------------------------------------------------------------- +-- e-class page wiring + +-- | Read every (eid, page blob) currently stored in the page table. +allPages :: SqlBackend db => PageStore db -> IO [(Int, BS.ByteString)] +allPages ps = do + rows <- queryDb (psDb ps) ("SELECT key, blob FROM " <> psTable ps) [] + pure [ (sqlToInt k, unhex (sqlToText b)) | [k, b] <- rows ] + +-- | Read every e-class id currently stored in the page table (keys only). This +-- does NOT load the page blobs, so callers that only need the id set (e.g. +-- 'recalculateBestAllStream' via 'cpsKeys') don't pull the whole page store +-- into memory just to extract keys. +allPageKeys :: SqlBackend db => PageStore db -> IO [Int] +allPageKeys ps = do + rows <- queryDb (psDb ps) ("SELECT key FROM " <> psTable ps) [] + pure [ sqlToInt k | [k] <- rows ] + +-- | Table used for the lazily paged e-class store (shared by save/load). +classStoreTable :: T.Text +classStoreTable = "cstore_page" + +-- | Table holding the re-saturation frontier: e-classes that have been created +-- or merged since the last frontier re-saturation pass, so a subsequent pass can +-- re-saturate only these (and classes they touch) instead of the whole graph. +frontierTable :: T.Text +frontierTable = "frontier" + +-- | Mark an e-class as recently changed (part of the re-saturation frontier). +markFrontier :: PageStore db -> Int -> IO () +markFrontier ps eid = modifyIORef' (psFrontier ps) (IntSet.insert eid) + +-- | Read the persisted frontier e-class ids (for initialising a pass). +loadFrontierRows :: SqlBackend db => db -> IO [Int] +loadFrontierRows db = do + rows <- queryDb db ("SELECT eid FROM " <> frontierTable) [] + pure [ sqlToInt e | [e] <- rows ] + +-- | Whether the matcher's candidate-root enumeration is restricted to the +-- frontier (a frontier re-saturation pass). Off by default (normal eqsat and +-- the pure in-memory path are unaffected). +setFrontierActive :: PageStore db -> Bool -> IO () +setFrontierActive ps a = writeIORef (psFrontierActive ps) a + +-- | Start a frontier re-saturation: seed the in-memory dirty set from the +-- persisted frontier (or an explicit seed), and restrict the matcher to it. +initFrontier :: SqlBackend db => PageStore db -> [Int] -> IO () +initFrontier ps seed = do + persisted <- loadFrontierRows (psDb ps) + writeIORef (psFrontier ps) (IntSet.fromList (persisted ++ seed)) + writeIORef (psFrontierActive ps) True + +-- | End a frontier re-saturation: clear the persisted frontier and the +-- in-memory dirty set, and lift the matcher restriction. +clearFrontier :: SqlBackend db => PageStore db -> IO () +clearFrontier ps = do + runDb (psDb ps) ("DELETE FROM " <> frontierTable) [] + writeIORef (psFrontier ps) IntSet.empty + writeIORef (psFrontierActive ps) False + +-- | Persist the in-memory frontier (recently-changed classes) to the frontier +-- table, so it survives until the next re-saturation pass. Idempotent. +flushFrontier :: SqlBackend db => PageStore db -> IO () +flushFrontier ps = do + f <- readIORef (psFrontier ps) + unless (IntSet.null f) $ do + let db = psDb ps + execDb db "BEGIN" + forM_ (IntSet.toList f) $ \eid -> + insertIgnore db ("frontier (eid, updated_at) VALUES (?, ?)") + [ SqlInteger (fromIntegral eid), SqlText (T.pack (show (0 :: Int))) ] + execDb db "COMMIT" + writeIORef (psFrontier ps) IntSet.empty + +-- | Open the e-class page store on 'classStoreTable', creating the table if +-- missing. @cap@ bounds the LRU cache (0 = unbounded), @flushEvery'@ is the +-- dirty-page threshold that triggers a batched write-back. +openClassStore :: SqlBackend db => db -> Int -> Int -> IO (PageStore db) +openClassStore db cap flushEvery' = newPageStore db classStoreTable cap flushEvery' + +-- | Adapt a 'PageStore' into the 'EClassPageStore' handle an 'EGraph' carries: +-- e-class blobs are 'Binary'-serialized pages. The store is authoritative; +-- the graph's resident map mirrors insertions and is consulted first on reads. +classStoreHandle :: SqlBackend db => PageStore db -> EClassPageStore +classStoreHandle ps = EClassPageStore + { cpsLookup = \eid -> fmap (fmap (decode . BL.fromStrict)) (readPage ps eid) + , cpsInsert = \ec -> writePage ps (_eClassId ec) (BL.toStrict (encode ec)) + , cpsDelete = \eid -> deletePage ps eid + , cpsFlush = writeback ps >> flushFrontier ps >> pure () + , cpsAll = fmap (map (decode . BL.fromStrict . snd)) (allPages ps) + , cpsKeys = allPageKeys ps + , cpsStreamRoots = \op budget exclude -> do + -- Union the DB's operator index with any not-yet-flushed nodes, so the + -- streaming matcher sees every live node (like the resident _patDB trie), + -- not just the last flushed snapshot. Excluded (already-attempted) roots + -- are skipped so the per-rule budget advances to new roots. When a + -- frontier re-saturation is active, roots are further restricted to the + -- frontier (recently-changed) e-classes. Memory is O(budget + pending + + -- size of exclude + size of frontier). + dbRoots <- streamByOp (psDb ps) (T.pack (opDetailOf op)) budget exclude + pend <- readIORef (psNodes ps) + active <- readIORef (psFrontierActive ps) + frontier <- readIORef (psFrontier ps) + let detail = opDetailOf op + ex = IntSet.fromList exclude + pendRoots = [ eid | (eid, en) <- Set.toList pend + , enodeOpDetail en == detail + , not (IntSet.member eid ex) ] + keep e = not active || IntSet.member e frontier + pure (take budget (filter keep (nub (pendRoots ++ dbRoots)))) + , cpsRecordNode = \en eid -> markFrontier ps eid >> modifyIORef' (psNodes ps) (Set.insert (eid, en)) + , cpsNodeToClass = \en -> do + -- content-address lookup, seeing not-yet-flushed nodes first (live) + pend <- readIORef (psNodes ps) + case [ eid | (eid, pe) <- Set.toList pend, pe == en ] of + (eid : _) -> pure (Just eid) + [] -> do + rows <- queryDb (psDb ps) + "SELECT eid FROM eclass_node WHERE enode_key = ?" + [SqlText (T.pack (enodeKey en))] + pure (case rows of [[v]] -> Just (sqlToInt v); _ -> Nothing) + , cpsCanonicalOf = \eid -> do + pend <- readIORef (psCanon ps) + case IM.lookup eid pend of + Just c -> pure (Just c) + Nothing -> do + rows <- queryDb (psDb ps) + "SELECT canonical FROM eclass WHERE eid = ?" + [SqlInteger (fromIntegral eid)] + pure (case rows of [[v]] -> Just (sqlToInt v); _ -> Nothing) + , cpsRecordCanonical = \eid canon -> markFrontier ps eid >> modifyIORef' (psCanon ps) (IM.insert eid canon) + , cpsBeginFrontier = initFrontier ps [] + , cpsEndFrontier = clearFrontier ps + } diff --git a/src/Algorithm/EqSat/Storage/Import.hs b/src/Algorithm/EqSat/Storage/Import.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Import.hs @@ -0,0 +1,312 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} + +-- | Out-of-core seed import: build an srtree e-graph directly in the database +-- by streaming a list of expressions into the relational schema. +-- +-- The import holds **no** graph-size data in RAM: e-nodes are content-addressed +-- against the @enode@/@eclass_node@ tables (the @enode_key@ is their content +-- address), child lookups for n-ary flattening read each child's node from the +-- DB, and parent edges are written straight to the @parent@ table. Only a +-- scalar next-id counter is kept in memory, so peak memory is bounded (one +-- e-class at a time) regardless of how many expressions are imported. +-- +-- Class pages (@cstore_page@) are written at the end in a single linear pass by +-- reconstructing each class from the relational tables, so hot classes are +-- never rewritten repeatedly. +-- +-- The produced database is byte-compatible with 'saveGraph': the same page +-- blobs and relational rows, so 'loadGraphLazy' / 'dbEqSat' work on it +-- unchanged. Saturation is deliberately NOT performed here (structural-only); +-- rule rewrites are left to the out-of-core 'dbEqSat' path. +module Algorithm.EqSat.Storage.Import + ( ImportSummary(..) + , importEqs + , recordExpressionIndex + ) where + +import Control.Monad (forM, forM_, foldM) +import Control.Exception (SomeException, catch, displayException, try) +import Data.IORef (IORef, newIORef, readIORef, modifyIORef', writeIORef) +import Data.Maybe (catMaybes, fromMaybe) +import qualified Data.Text as T + +import qualified Data.IntMap as IntMap +import qualified Data.HashSet as HashSet + +import Data.Binary (encode) +import qualified Data.ByteString.Lazy as BL + +import Data.SRTree +import Data.SRTree.Eval (Target) + +import Algorithm.EqSat.Egraph + ( EClassId, ENode(..), NOp(..), EClass(..), EClassData(..), Consts(..) ) +import Algorithm.EqSat.Storage.Backend + ( SqlValue(..), SqlBackend(..), sqlToInt, sqlToMaybeDouble, sqlToText ) +import Algorithm.EqSat.Storage.ClassStore (classStoreTable, hex) +import Algorithm.EqSat.Storage.Types + ( enodeKey, enodeOpTag, enodeOpDetail, serializeTheta, parseTheta, parseEnodeKey ) +import Algorithm.EqSat.Storage.Schema (createSchema) +import Algorithm.EqSat.Storage.Query (getOrCreateDataset, writeDatasetFit) + +-- | Result of an out-of-core import. +data ImportSummary = ImportSummary + { isNextId :: !Int -- ^ next free e-class id + , isClasses :: !Int -- ^ total e-classes written + , isExpressions :: !Int -- ^ root expressions inserted + } deriving (Show, Eq) + +-- | The only in-process state: the next free e-class id (a single scalar, so +-- import memory is independent of the number of classes). +data ImportState = ImportState + { stNextId :: !Int + } + +-- | Insert a list of @(expression, theta, fitness)@ into the database, +-- 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 + 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) + r <- try $ do + execDb db "BEGIN" + -- stream the expression list through a fold (rather than forM_ + length + -- eqs) so the lazy list is unreferenced after consumption and GC'd as it + -- 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) ] + pure (c + 1)) 0 eqs + writeMeta db ref + writeAllPages db + execDb db "COMMIT" + pure n + case r of + Left (e :: SomeException) -> do + _ <- (execDb db "ROLLBACK" `catch` \(_ :: SomeException) -> pure ()) + pure (Left ("importEqs failed: " <> displayException e)) + Right n -> do + st <- readIORef ref + 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) [] + 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 + 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)] + +-- | 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 + 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 + +-- | 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 +-- from the database (no in-memory class index). +flattenChildren :: SqlBackend db => db -> NOp -> [(EClassId, Int)] -> IO (IntMap.IntMap Int) +flattenChildren db op children = do + ms <- forM children $ \(c, n) -> do + men <- lookupClassNode db c + case men of + Just (ENAry op' m') | op' == op -> pure (IntMap.map (* n) m') + _ -> pure (IntMap.singleton c n) + pure (IntMap.unionsWith (+) ms) + +-- | 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 + let key = enodeKey en + childs = dedupChildren children + mEid <- lookupEnodeId db key + case mEid of + Just eid -> do + h <- classHeight db eid + pure (eid, h) + Nothing -> do + st <- readIORef ref + 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 + pure (eid, h) + +-- | Merge duplicate child e-classes into multiplicities (e.g. @x0 - x0@ has +-- both children in the same class), keeping the max height. +dedupChildren :: [(EClassId, Int, Int)] -> [(EClassId, Int, Int)] +dedupChildren = + map (\(c, (n, h)) -> (c, n, h)) . IntMap.toList + . IntMap.fromListWith (\(n1, h1) (n2, h2) -> (n1 + n2, max h1 h2)) + . map (\(c, n, h) -> (c, (n, h))) + +-- | 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 + runDb db "INSERT INTO enode (key, op, op_detail) VALUES (?, ?, ?)" + [ SqlText (T.pack key) + , SqlText (T.pack (enodeOpTag en)) + , SqlText (T.pack (enodeOpDetail en)) ] + runDb db "INSERT INTO eclass (eid, canonical, height) VALUES (?, ?, ?)" + [ SqlInteger (fromIntegral eid) + , SqlInteger (fromIntegral eid) + , SqlInteger (fromIntegral h) ] + runDb db "INSERT INTO eclass_node (eid, enode_key) VALUES (?, ?)" + [ SqlInteger (fromIntegral eid), SqlText (T.pack key) ] + -- enode_child rows exist only for ENAry nodes (EBin/Uni children live in the + -- content key); the multiset is unique. + forM_ (naryChildrenOf en) $ \(c, n) -> + runDb db "INSERT INTO enode_child (enode_key, child_eid, cnt) VALUES (?, ?, ?)" + [ SqlText (T.pack key) + , SqlInteger (fromIntegral c) + , SqlInteger (fromIntegral n) ] + forM_ children $ \(c, _, _) -> + runDb db "INSERT INTO parent (child_eid, parent_eid, parent_enode_key) VALUES (?, ?, ?)" + [ SqlInteger (fromIntegral c) + , SqlInteger (fromIntegral eid) + , 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 + +-- | ENAry children as (class, multiplicity); empty for all other node shapes +-- (their children live inline in the content key). +naryChildrenOf :: ENode -> [(EClassId, Int)] +naryChildrenOf (ENAry _ m) = IntMap.toList m +naryChildrenOf _ = [] + +-- | Look up an existing e-class id for a content key (NULL if absent). +lookupEnodeId :: SqlBackend db => db -> String -> IO (Maybe EClassId) +lookupEnodeId db key = do + rows <- queryDb db "SELECT eid FROM eclass_node WHERE enode_key = ?" + [SqlText (T.pack key)] + pure $ case rows of + ([eid] : _) -> Just (sqlToInt eid) + _ -> Nothing + +-- | Height of an e-class (from the @eclass@ row). +classHeight :: SqlBackend db => db -> EClassId -> IO Int +classHeight db eid = do + rows <- queryDb db "SELECT height FROM eclass WHERE eid = ?" + [SqlInteger (fromIntegral eid)] + pure $ case rows of + ([h] : _) -> sqlToInt h + _ -> 0 + +-- | The singleton e-node of a class (NULL if the class has no node row yet). +lookupClassNode :: SqlBackend db => db -> EClassId -> IO (Maybe ENode) +lookupClassNode db eid = do + rows <- queryDb db "SELECT enode_key FROM eclass_node WHERE eid = ?" + [SqlInteger (fromIntegral eid)] + pure $ case rows of + ([SqlText k] : _) -> parseEnodeKey (T.unpack k) + _ -> Nothing + +-- | Record that an expression (by its canonical root e-node) was seen in a +-- dataset's graph, so "was this expression already tested?" is answerable per +-- dataset. Used by the delta-insert path ('dbInsert') to keep @expression_index@ +-- live for newly-added expressions. +recordExpressionIndex :: SqlBackend db => db -> Int -> EClassId -> IO () +recordExpressionIndex db dsid eid = do + 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) ] + +-- | Write every class page once, reconstructing each class from the relational +-- tables (node + height from @eclass_node@/@eclass@, parents from @parent@, +-- metrics from @fit@). Memory stays bounded to a single class at a time. +writeAllPages :: SqlBackend db => db -> IO () +writeAllPages db = do + cids <- queryDb db "SELECT eid FROM eclass ORDER BY eid" [] + forM_ cids $ \[eidCol] -> do + let eid = sqlToInt eidCol + nh <- queryDb db + "SELECT n.enode_key, c.height FROM eclass_node n \ + \JOIN eclass c ON c.eid = n.eid WHERE n.eid = ?" + [SqlInteger (fromIntegral eid)] + case nh of + ([SqlText k, hcol] : _) -> do + let en = fromMaybe (error ("importEqs: bad node key for eid " <> show eid)) + (parseEnodeKey (T.unpack k)) + h = sqlToInt hcol + pr <- queryDb db "SELECT parent_eid, parent_enode_key FROM parent WHERE child_eid = ?" + [SqlInteger (fromIntegral eid)] + let parents = HashSet.fromList + [ (sqlToInt pe, fromMaybe (error "importEqs: bad parent key") (parseEnodeKey (T.unpack (sqlToText pk)))) + | [pe, pk] <- pr ] + -- Pages are structural-only: cost/best are derived on load, and + -- fitness/dl/theta are dataset metadata (dataset_fit), so they are NOT + -- baked into the e-graph blob (keeps the graph reusable across datasets). + writeClassPage db eid (EClass eid (HashSet.singleton en) parents h (defaultInfo en h)) + _ -> pure () + +-- | Serialize an e-class to the page store (INSERT OR REPLACE so the final +-- pass is idempotent). +writeClassPage :: SqlBackend db => db -> EClassId -> EClass -> IO () +writeClassPage db eid ec = + runDb db ("INSERT OR REPLACE INTO " <> classStoreTable <> " (key, blob) VALUES (?, ?)") + [ SqlInteger (fromIntegral eid) + , SqlText (hex (BL.toStrict (encode ec))) ] + +-- | Per-class data. Cost/best are derived quantities recomputed on load +-- ('recalculateBestAll'). Fitness/dl/theta/size are baked into the page (so +-- out-of-core reads via the paged store see them). @_consts@ is set from the +-- node so constant-folding rules behave identically to the in-memory seed. +defaultInfo :: ENode -> Int -> EClassData +defaultInfo en h = EData 0 en (constOf en) Nothing Nothing [] h + +constOf :: ENode -> Consts +constOf (EConst x) = ConstVal x +constOf (EParam ix) = ParamIx ix +constOf _ = NotConst + +-- | Write the @meta@ scalars: next free id and DB-tracking flag (mirrors the +-- in-memory seed graph, which runs with range-DB tracking enabled). +writeMeta :: SqlBackend db => db -> IORef ImportState -> IO () +writeMeta db ref = do + st <- readIORef ref + runDb db "INSERT INTO meta (key, value) VALUES (?, ?)" + [ SqlText "next_id", SqlText (T.pack (show (stNextId st))) ] + runDb db "INSERT INTO meta (key, value) VALUES (?, ?)" + [ SqlText "track_dbs", SqlText "1" ] diff --git a/src/Algorithm/EqSat/Storage/Postgres.hs b/src/Algorithm/EqSat/Storage/Postgres.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Postgres.hs @@ -0,0 +1,213 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | PostgreSQL-backed persistence for srtree e-graphs. +-- +-- Implements the same driver-neutral interface as +-- 'Algorithm.EqSat.Storage.SQLite' on top of @libpq@ +-- ('Database.PostgreSQL.LibPQ'), so the shared storage code +-- ('saveGraph'/'loadGraph'/'pushFit'/'refreshFitness' and the +-- 'Algorithm.EqSat.Storage.Query' API) runs unchanged against PostgreSQL. +-- +-- Connections are plain @libpq@ connections (see 'connectPostgres' / +-- 'closePostgres'); the reggression layer dispatches on a @postgres://@ / +-- @postgresql://@ DSN. +module Algorithm.EqSat.Storage.Postgres + ( schemaPostgres + , connectPostgres + , closePostgres + ) where + +import Control.Monad (forM, forM_) +import Data.ByteString (ByteString) +import qualified Data.IntSet as IntSet +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import Database.PostgreSQL.LibPQ + ( Connection, ExecStatus(..), Format(..), Oid(..), Result + , connectdb, exec, execParams, finish, getvalue, invalidOid, nfields + , ntuples, resultErrorMessage, resultStatus, toColumn, toRow ) + +import Algorithm.EqSat.Storage.Backend (SqlValue(..), SqlBackend(..), sqlToInt, sqlToText) +import Algorithm.EqSat.Storage.ClassStore (unhex) + +-- | PostgreSQL DDL. Mirrors 'Algorithm.EqSat.Storage.Schema.schemaSQL'. +-- +-- Differences from SQLite: @BIGINT@ identity keys, @DOUBLE PRECISION@ +-- metrics, and foreign keys declared @DEFERRABLE INITIALLY DEFERRED@ so the +-- writer can insert @eclass_node@/@fit@ rows before their referenced +-- @eclass@ rows within the @BEGIN@..@COMMIT@ transaction of 'saveGraph'. +schemaPostgres :: [Text] +schemaPostgres = + [ "CREATE TABLE IF NOT EXISTS meta (" + <> " key TEXT PRIMARY KEY," + <> " value TEXT NOT NULL)" + , "CREATE TABLE IF NOT EXISTS enode (" + <> " key TEXT PRIMARY KEY," + <> " op TEXT NOT NULL," + <> " op_detail TEXT," + <> " a BIGINT," + <> " b BIGINT," + <> " x DOUBLE PRECISION)" + , "CREATE TABLE IF NOT EXISTS enode_child (" + <> " enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " child_eid BIGINT NOT NULL," + <> " cnt INTEGER NOT NULL DEFAULT 1," + <> " PRIMARY KEY (enode_key, child_eid))" + , "CREATE TABLE IF NOT EXISTS eclass (" + <> " eid BIGINT PRIMARY KEY," + <> " canonical BIGINT NOT NULL," + <> " height INTEGER NOT NULL DEFAULT 0)" + , "CREATE TABLE IF NOT EXISTS eclass_node (" + <> " eid BIGINT NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " PRIMARY KEY (eid, enode_key))" + , "CREATE TABLE IF NOT EXISTS parent (" + <> " child_eid BIGINT NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " parent_eid BIGINT NOT NULL," + <> " parent_enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " PRIMARY KEY (child_eid, parent_eid, parent_enode_key))" + , "CREATE TABLE IF NOT EXISTS cstore_page (" + <> " key TEXT PRIMARY KEY," + <> " blob TEXT NOT NULL)" + , "CREATE TABLE IF NOT EXISTS frontier (" + <> " eid BIGINT PRIMARY KEY REFERENCES eclass(eid) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED," + <> " updated_at TEXT)" + , "CREATE TABLE IF NOT EXISTS dataset (" + <> " id BIGSERIAL PRIMARY KEY," + <> " name TEXT NOT NULL UNIQUE," + <> " created TEXT)" + , "CREATE TABLE IF NOT EXISTS dataset_fit (" + <> " dataset_id BIGINT NOT NULL REFERENCES dataset(id) ON DELETE CASCADE," + <> " eid BIGINT NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " fitness DOUBLE PRECISION," + <> " dl DOUBLE PRECISION," + <> " theta TEXT," + <> " size INTEGER NOT NULL DEFAULT 0," + <> " evaluated INTEGER NOT NULL DEFAULT 0," + <> " fitted INTEGER NOT NULL DEFAULT 0," + <> " stale INTEGER NOT NULL DEFAULT 0," + <> " updated_at TEXT," + <> " PRIMARY KEY (dataset_id, eid))" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_fitness ON dataset_fit(fitness)" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_size ON dataset_fit(size)" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_dl ON dataset_fit(dl)" + , "CREATE TABLE IF NOT EXISTS expression_index (" + <> " expression_key TEXT PRIMARY KEY," + <> " eclass BIGINT NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " dataset_id BIGINT REFERENCES dataset(id) ON DELETE CASCADE," + <> " first_seen TEXT)" + ] + +-- | Open a PostgreSQL connection from a connection string (e.g. +-- @postgresql://user:pass@host:5432/db@). +connectPostgres :: String -> IO Connection +connectPostgres = connectdb . TE.encodeUtf8 . T.pack + +-- | Close a PostgreSQL connection. +closePostgres :: Connection -> IO () +closePostgres = finish + +instance SqlBackend Connection where + execDb conn sql = do + r <- pgExec conn sql + statusOK r "exec" + + runDb conn sql params = do + r <- pgExecParams conn sql params + statusOK r "run" + + insertIgnore conn tail params = do + r <- pgExecParams conn ("INSERT INTO " <> tail <> " ON CONFLICT DO NOTHING") params + statusOK r "insertIgnore" + + queryDb conn sql params = do + r <- pgExecParams conn sql params + st <- resultStatus r + case st of + TuplesOk -> do + ns <- ntuples r + nf <- nfields r + let n = fromEnum ns + m = fromEnum nf + forM [0 .. n - 1] $ \i -> + forM [0 .. m - 1] $ \j -> do + v <- getvalue r (toRow i) (toColumn j) + pure $ case v of + Nothing -> SqlNull + Just bs -> SqlText (TE.decodeUtf8 bs) + _ -> do + statusOK r "query" + pure [] + + createSchemaDb conn = mapM_ (execDb conn) schemaPostgres + + -- Grid fallback (Postgres is not the out-of-core target): the cursor-based + -- streaming matcher needs 'Database.SQLite3'; here we return the full + -- distinct set up to @budget@, documented as unbounded memory. + streamByOp conn opDetail budget exclude = do + let ex = IntSet.fromList exclude + rows <- queryDb conn + "SELECT DISTINCT n.eid FROM eclass_node n \ + \JOIN enode e ON e.key = n.enode_key WHERE e.op_detail = ?" + [SqlText opDetail] + pure (take budget [ eid | [eid'] <- rows, let eid = sqlToInt eid', not (IntSet.member eid ex) ]) + -- Grid fallback for page streaming (unbounded; Postgres is not the + -- out-of-core target). + streamPages conn tbl k = do + rows <- queryDb conn ("SELECT key, blob FROM " <> tbl) [] + forM_ rows $ \[key, blob] -> k (fromIntegral (sqlToInt key)) (unhex (sqlToText blob)) + +-- | Raise an exception unless the status is @CommandOk@/@TuplesOk@. +statusOK :: Result -> Text -> IO () +statusOK r tag = do + st <- resultStatus r + case st of + CommandOk -> pure () + TuplesOk -> pure () + EmptyQuery -> pure () + _ -> do + mmsg <- resultErrorMessage r + let msg = maybe "unknown error" (T.unpack . TE.decodeUtf8) mmsg + fail ("postgres: " <> T.unpack tag <> ": " <> msg) + +-- | Execute a statement without parameters (DDL, BEGIN/COMMIT, DELETE). +pgExec :: Connection -> Text -> IO Result +pgExec conn sql = do + mr <- exec conn (TE.encodeUtf8 sql) + case mr of + Nothing -> fail "postgres: exec returned no result" + Just r -> pure r + +-- | Execute a parameterized statement, rewriting @?@ to @$n@. +pgExecParams :: Connection -> Text -> [SqlValue] -> IO Result +pgExecParams conn sql params = do + let pgSql = toPG sql + ps = map renderParam params + mr <- execParams conn (TE.encodeUtf8 pgSql) ps Text + case mr of + Nothing -> fail "postgres: execParams returned no result" + Just r -> pure r + +-- | Render a parameter for libpq's text-format protocol. NULL is never sent +-- as a parameter: the shared SQL spells it as the literal @NULL@. +renderParam :: SqlValue -> Maybe (Oid, ByteString, Format) +renderParam (SqlInteger n) = Just (invalidOid, TE.encodeUtf8 (T.pack (show n)), Text) +renderParam (SqlFloat d) = Just (invalidOid, TE.encodeUtf8 (T.pack (show d)), Text) +renderParam (SqlText t) = Just (invalidOid, TE.encodeUtf8 t, Text) +renderParam SqlNull = Nothing + +-- | Rewrite the shared positional @?@ placeholders to libpq's @$n@ form +-- (single-quoted literals are skipped, in case a value embeds @?@). +toPG :: Text -> Text +toPG = T.pack . go (1 :: Int) . T.unpack + where + go :: Int -> String -> String + go _ [] = [] + go n ('\'' : r) = '\'' : skip r + where + skip ('\'' : r') = '\'' : go n r' + skip (c : r') = c : skip r' + skip [] = [] + go n ('?' : r) = '$' : show n ++ go (n + 1) r + go n (c : r) = c : go n r diff --git a/src/Algorithm/EqSat/Storage/Query.hs b/src/Algorithm/EqSat/Storage/Query.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Query.hs @@ -0,0 +1,189 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | SQL query API over a stored e-graph's @fit@ / @enode@ tables. +-- +-- The queries are intentionally SQL-shaped (this is the slice of the +-- reggression functionality that runs directly in the database) and mirror the +-- in-memory counterparts in 'Algorithm.EqSat.Queries'. They are written +-- against 'Algorithm.EqSat.Storage.Backend', so the same SQL drives the +-- SQLite and PostgreSQL backends. +module Algorithm.EqSat.Storage.Query + ( getOrCreateDataset + , datasetId + , firstDatasetId + , writeDatasetFit + , readDatasetFit + , topN + , pareto + , paretoBySize + , distributionCounts + , countPattern + , expressionEclass + , testedOnDataset + , versionsOf + ) where + +import Data.Maybe (catMaybes) +import Data.Text (Text) +import qualified Data.Text as T + +import Algorithm.EqSat.Egraph (EClassId) + +import Algorithm.EqSat.Storage.Backend + ( SqlBackend, SqlValue(..), runDb, queryDb, sqlToInt, sqlToMaybeDouble, sqlToText ) +import Algorithm.EqSat.Storage.Schema (createSchema) + +-- | Resolve a dataset name to its @dataset@ row id, creating it if needed. +getOrCreateDataset :: SqlBackend db => db -> String -> IO Int +getOrCreateDataset db name = do + createSchema db + m <- datasetId db name + case m of + Just i -> pure i + Nothing -> do + runDb db "INSERT INTO dataset (name) VALUES (?)" [SqlText (T.pack name)] + r <- datasetId db name + pure (maybe 0 id r) + +-- | Look up an existing dataset id by name. +datasetId :: SqlBackend db => db -> String -> IO (Maybe Int) +datasetId db name = do + rows <- queryDb db "SELECT id FROM dataset WHERE name = ?" [SqlText (T.pack name)] + pure $ case rows of + ([i] : _) -> Just (sqlToInt i) + _ -> Nothing + +-- | The id of the first dataset (used by the legacy 'loadGraph' path, which is +-- not dataset-scoped, to source per-class risk metrics from @dataset_fit@). +firstDatasetId :: SqlBackend db => db -> IO (Maybe Int) +firstDatasetId db = do + rows <- queryDb db "SELECT id FROM dataset ORDER BY id LIMIT 1" [] + pure $ case rows of + ([i] : _) -> Just (sqlToInt i) + _ -> Nothing + +-- | Upsert a per-(dataset, e-class) fit row. +writeDatasetFit + :: SqlBackend db => db -> Int -> EClassId + -> Maybe Double -> Maybe Double -> Text -> Int -> IO () +writeDatasetFit db ds eid fit dl theta sz = do + let (fitCol, fitVal) = case fit of + Nothing -> ("NULL", Nothing) + Just f -> ("?", Just (SqlFloat f)) + (dlCol, dlVal) = case dl of + Nothing -> ("NULL", Nothing) + Just d -> ("?", Just (SqlFloat d)) + runDb db + ("INSERT OR REPLACE INTO dataset_fit \ + \(dataset_id, eid, fitness, dl, theta, size, evaluated, fitted) \ + \VALUES (?, ?, " <> fitCol <> ", " <> dlCol <> ", ?, ?, 1, 1)") + (catMaybes [ Just (SqlInteger (fromIntegral ds)) + , Just (SqlInteger (fromIntegral eid)) + , fitVal + , dlVal + , Just (SqlText theta) + , Just (SqlInteger (fromIntegral sz)) ]) + +-- | Read per-(dataset, e-class) fit rows. +readDatasetFit + :: SqlBackend db => db -> Int + -> IO [(EClassId, (Maybe Double, Maybe Double, Int, Text))] +readDatasetFit db ds = do + createSchema db + rows <- queryDb db + "SELECT eid, fitness, dl, size, theta FROM dataset_fit WHERE dataset_id = ?" + [SqlInteger (fromIntegral ds)] + pure [ (sqlToInt eid, (sqlToMaybeDouble f, sqlToMaybeDouble d, sqlToInt sz, sqlToText th)) + | [eid, f, d, sz, th] <- rows ] + +-- | The @n@ e-classes with the best fitness (on the dataset), descending. +topN :: SqlBackend db => db -> Int -> Int -> IO [(EClassId, Double)] +topN db ds n = do + rows <- queryDb db + "SELECT eid, fitness FROM dataset_fit \ + \WHERE dataset_id = ? AND fitness IS NOT NULL \ + \ORDER BY fitness DESC LIMIT ?" + [ SqlInteger (fromIntegral ds), SqlInteger (fromIntegral n) ] + pure [ (sqlToInt eid, f) + | [eid, f] <- rows + , Just f <- [sqlToMaybeDouble f] ] + +-- | Non-dominated classes over (max fitness, min dl) on the dataset. Returns +-- the (eid, fitness, dl) triples that are not dominated by any other class. +pareto :: SqlBackend db => db -> Int -> IO [(EClassId, Double, Double)] +pareto db ds = do + rows <- queryDb db + "SELECT eid, fitness, dl FROM dataset_fit \ + \WHERE dataset_id = ? AND fitness IS NOT NULL AND dl IS NOT NULL" + [SqlInteger (fromIntegral ds)] + let pts = [ (sqlToInt eid, f, d) + | [eid, ff, dd] <- rows + , Just f <- [sqlToMaybeDouble ff] + , Just d <- [sqlToMaybeDouble dd] ] + dominates (f1, d1) (f0, d0) = f1 >= f0 && d1 <= d0 && (f1 > f0 || d1 < d0) + nonDominated (eid_, f, d) = not (any (\(q0, qf, qd) -> dominates (qf, qd) (f, d)) pts) + pure [ p | p@(_, f, d) <- pts, nonDominated p ] + +-- | Non-dominated classes over (max fitness, min size) on the dataset. +paretoBySize :: SqlBackend db => db -> Int -> IO [(EClassId, Double, Int)] +paretoBySize db ds = do + rows <- queryDb db + "SELECT eid, fitness, size FROM dataset_fit \ + \WHERE dataset_id = ? AND fitness IS NOT NULL" + [SqlInteger (fromIntegral ds)] + let pts = [ (sqlToInt eid, f, s) + | [eid, ff, ss] <- rows + , Just f <- [sqlToMaybeDouble ff] + , let s = sqlToInt ss ] + dominates (f1, s1) (f0, s0) = f1 >= f0 && s1 <= s0 && (f1 > f0 || s1 < s0) + nonDominated (eid_, f, s) = not (any (\(q0, qf, qs) -> dominates (qf, qs) (f, s)) pts) + pure [ p | p@(_, f, s) <- pts, nonDominated p ] + +-- | Number of evaluated e-classes per model size (up to @maxSize@) on the +-- dataset. +distributionCounts :: SqlBackend db => db -> Int -> Int -> IO [(Int, Int)] +distributionCounts db ds maxSize = do + rows <- queryDb db + "SELECT size, COUNT(*) FROM dataset_fit \ + \WHERE dataset_id = ? AND fitness IS NOT NULL AND size <= ? \ + \GROUP BY size ORDER BY size" + [ SqlInteger (fromIntegral ds), SqlInteger (fromIntegral maxSize) ] + pure [ (sqlToInt s, sqlToInt c) | [s, c] <- rows ] + +-- | The e-class a previously-indexed expression maps to (NULL if never seen). +expressionEclass :: SqlBackend db => db -> Text -> IO (Maybe EClassId) +expressionEclass db key = do + rows <- queryDb db "SELECT eclass FROM expression_index WHERE expression_key = ?" + [SqlText key] + pure $ case rows of + ([e] : _) -> Just (sqlToInt e) + _ -> Nothing + +-- | Whether a class has a fitness (i.e. was evaluated/fitted) on a dataset. +testedOnDataset :: SqlBackend db => db -> Int -> EClassId -> IO Bool +testedOnDataset db ds eid = do + rows <- queryDb db + "SELECT 1 FROM dataset_fit WHERE dataset_id = ? AND eid = ? AND fitness IS NOT NULL" + [SqlInteger (fromIntegral ds), SqlInteger (fromIntegral eid)] + pure (not (null rows)) + +-- | The e-node content keys that make up an e-class (multiple versions of one +-- expression). +versionsOf :: SqlBackend db => db -> EClassId -> IO [Text] +versionsOf db eid = do + rows <- queryDb db "SELECT enode_key FROM eclass_node WHERE eid = ?" + [SqlInteger (fromIntegral eid)] + pure [ sqlToText k | [k] <- rows ] + +-- | Number of distinct e-classes containing at least one e-node whose +-- specific operator matches (e.g. \"EAdd\", \"EMul\", \"Add\", \"LogAbs\"). +countPattern :: SqlBackend db => db -> Text -> IO Int +countPattern db op = do + rows <- queryDb db + "SELECT COUNT(DISTINCT eclass_node.eid) \ + \FROM eclass_node JOIN enode ON enode.key = eclass_node.enode_key \ + \WHERE enode.op_detail = ?" + [ SqlText op ] + pure $ case rows of + row : _ | [n] <- row -> sqlToInt n + _ -> 0 diff --git a/src/Algorithm/EqSat/Storage/SQLite.hs b/src/Algorithm/EqSat/Storage/SQLite.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/SQLite.hs @@ -0,0 +1,592 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | 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 + , loadGraph + , loadGraphLazy + , pushFit + , refreshFitness + , query + , flushStore + ) where + +import Control.Monad (forM, forM_, when) +import Control.Exception (SomeException, catch, displayException) +import Control.Monad.Identity (runIdentity) +import Control.Monad.State.Strict (execStateT) +import Data.Int (Int64) +import Data.Maybe (catMaybes, fromMaybe, listToMaybe) +import Data.Text (Text) +import qualified Data.Text as T +import Data.List (foldl') +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 Database.SQLite3 + ( Database, SQLData(..), StepResult(..) + , bind, columns, exec, step, withStatement ) + +import Data.SRTree.Eval (Target) +import Algorithm.EqSat.Egraph + ( EGraph(..), 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 ) +import Algorithm.EqSat.Storage.ClassStore + ( classStoreTable, openClassStore, allPages, classStoreHandle, hex, unhex ) +import Algorithm.EqSat.Storage.Query (readDatasetFit, writeDatasetFit, firstDatasetId) +import Algorithm.EqSat.Storage.Stream (streamRootsByOp) +import Algorithm.EqSat.Storage.Types +import Algorithm.EqSat.Storage.Schema (createSchema, schemaSQL) + +-- | Default cache capacity (pages) for the lazily paged e-class store. +defaultClassCap :: Int +defaultClassCap = 50000 + +-- --------------------------------------------------------------------------- +-- 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) + -- 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 + hv = case cols of (_ : SQLText t : _) -> t; _ -> "" + k eid (unhex hv) + go stmt + createSchemaDb db = mapM_ (exec db) schemaSQL + +toSqlData :: SqlValue -> SQLData +toSqlData (SqlInteger n) = SQLInteger n +toSqlData (SqlFloat d) = SQLFloat d +toSqlData (SqlText t) = SQLText t +toSqlData SqlNull = SQLNull + +fromSqlData :: SQLData -> SqlValue +fromSqlData (SQLInteger n) = SqlInteger n +fromSqlData (SQLFloat d) = SqlFloat d +fromSqlData (SQLText t) = SqlText t +fromSqlData SQLNull = SqlNull +fromSqlData _ = 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 = do + createSchema db + let rows0 = exportEGraph eg + case _classStore eg of + Just _ -> commitWith $ do + writeMeta db rows0 + Nothing -> do + gr <- graphClassRows eg + let rows = rows0 { _grEClasses = gr } + commitWith $ do + clearTables db + writeMeta db rows + writeNodes db rows + writeClasses db rows + writeParents db rows + writeDatasetFitRows db dsid rows + 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 + +-- | 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 parent" + 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) + , SqlText (hex (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 INTO meta (key, value) VALUES (?, ?)" + [ SqlText "next_id", SqlText (T.pack (show (_grNextId rows))) ] + run db "INSERT 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 + run db "INSERT INTO enode (key, op, op_detail) VALUES (?, ?, ?)" + [ SqlText (T.pack key) + , SqlText (T.pack (enodeOpTag en)) + , SqlText (T.pack (enodeOpDetail en)) ] + run db "INSERT INTO eclass_node (eid, enode_key) VALUES (?, ?)" + [ SqlInteger (fromIntegral eid), SqlText (T.pack key) ] + forM_ (naryChildren en) $ \(c, n) -> + run db "INSERT 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)))) ] + +-- | Persist the reverse edges: for every (parent class, parent e-node) in each +-- class's @_parents@, a @parent@ row keyed by the child e-class. This makes the +-- parent relation queryable per class without scanning @enode@/@eclass_node@. +writeParents :: SqlBackend db => db -> GraphRows -> IO () +writeParents db rows = + forM_ (IntMap.toAscList (_grEClasses rows)) $ \(eid, r) -> + forM_ (Set.toList (_rcParents r)) $ \(pEid, pEn) -> + run db "INSERT INTO parent (child_eid, parent_eid, parent_enode_key) VALUES (?, ?, ?)" + [ SqlInteger (fromIntegral eid) + , SqlInteger (fromIntegral pEid) + , SqlText (T.pack (enodeKey pEn)) ] + +-- | 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 + +-- | 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 + fit <- case mdsid of + Nothing -> pure [] + Just ds -> do + rows <- readDatasetFit db ds + pure [ (eid, (f, d, sz, parseTheta (T.unpack th))) + | (eid, (f, d, sz, th)) <- rows ] + 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) + parents <- readParents db + let storedParents = IntMap.fromListWith Set.union + [ (c, Set.singleton (pEid, pEn)) + | (c, pEid, pEn) <- parents ] + classes = buildClasses canon nodeToEClass storedParents (IntMap.fromList fit) (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 fitMap = IntMap.fromList fit + applyFit eid ec = + case IntMap.lookup eid fitMap of + Nothing -> ec + Just (f, d, s, th) -> + ec { _info = (_info ec){ _fitness = f, _dl = d, _size = s, _theta = th } } + classes = IntMap.mapWithKey applyFit + (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) }) + +-- | 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 + -> EGraphDB +seedEDBPaged nextId trackDBs _nodeToEClass = + (emptyDB){ _nextId = nextId, _trackDBs = trackDBs } + +-- | 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 -> IO (Either String EGraph) +loadGraphLazy db dsid = do + m <- readMeta db + case m of + Nothing -> pure (Left "srtree-db: no e-graph stored in this database") + Just (nextId, trackDBs) -> do + dsFit <- readDatasetFit db dsid + let fitMap = IntMap.fromList [ (eid, (f, d, sz, parseTheta (T.unpack th))) + | (eid, (f, d, sz, th)) <- dsFit ] + -- 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 + parents <- readParents db + let storedParents = IntMap.fromListWith Set.union + [ (c, Set.singleton (pEid, pEn)) + | (c, pEid, pEn) <- parents ] + classes = buildClasses canon0 nodeToEClass storedParents 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 fitSlim eid)) (cpsLookup base eid) } + eDB = seedEDBPaged nextId trackDBs HashMap.empty + eg = EGraph IntMap.empty HashMap.empty IntMap.empty eDB (Just h) + pure (Right 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 + :: IntMap.IntMap (Maybe Double, Maybe Double, Int) + -> EClassId -> EClass -> EClass +applyDsFit m eid ec = case IntMap.lookup eid m of + Nothing -> ec + Just (f, dl, sz) -> ec { _info = (_info ec) { _fitness = f, _dl = dl, _size = sz } } + +-- | 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 ] + +-- | Read (child e-class, parent e-class, parent e-node) edges from the @parent@ +-- table. The rows are grouped per child class by 'loadGraph'. +readParents :: SqlBackend db => db -> IO [(EClassId, EClassId, ENode)] +readParents db = do + rows <- query db "SELECT child_eid, parent_eid, parent_enode_key FROM parent" [] + pure (catMaybes + [ do + en <- parseEnodeKey (T.unpack (sqlToText k)) + pure (sqlToInt c, sqlToInt p, en) + | [c, p, k] <- rows ]) + +-- | Rebuild @_grEClasses@ rows: only canonical roots carry real class rows. +-- Parent pointers come from the stored @parent@ relation when present, falling +-- back to recomputation from the node -> class map (e.g. databases written +-- before the @parent@ table existed, or hand-built rows). +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 +-- (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 db dsid eg = do + createSchema db + run db "DELETE FROM dataset_fit WHERE dataset_id = ?" [SqlInteger (fromIntegral dsid)] + case _classStore eg of + Nothing -> do + let rows0 = exportEGraph eg + writeDatasetFitRows db dsid rows0 + Just _ -> do + -- flush so every dirty page is in @cstore_page@, then stream each class + -- and write its row, discarding it (O(1) memory). + flushStore eg + execDb db "BEGIN" + streamPages db classStoreTable $ \_eid blob -> do + let ec = decode (BL.fromStrict blob) + info = _info ec + writeDatasetFit db dsid (_eClassId ec) (_fitness info) (_dl info) + (T.pack (serializeTheta (_theta info))) (_size info) + execDb db "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 m = forM_ dsFit $ \(eid, (fitM, _, _, theta)) -> + case fitM of + Nothing -> pure () + Just f -> do + c <- canonical eid + insertFitness c f (parseTheta (T.unpack theta)) + pure (Right (runIdentity $ execStateT m eg)) diff --git a/src/Algorithm/EqSat/Storage/Schema.hs b/src/Algorithm/EqSat/Storage/Schema.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Schema.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Schema for persisting srtree e-graphs. +-- +-- Layout (shared by the SQLite and PostgreSQL backends): +-- * @meta@ - scalar settings (@next_id@, @track_dbs@, cost-function tag) +-- * @enode@ - content-addressable e-nodes (@key@ = canonical serialization) +-- * @enode_child@ - ENAry multiset children (@child_eid@, @cnt@) +-- * @eclass@ - e-class id -> canonical representative + height +-- * @eclass_node@ - canonical e-node -> e-class membership +-- * @parent@ - reverse edges: child e-class -> (parent e-class, parent e-node) +-- * @fit@ - per-class risk metrics (fitness, dl, size, theta) +-- +-- 'schemaSQL' is the SQLite DDL; 'Algorithm.EqSat.Storage.Postgres' carries +-- the equivalent PostgreSQL DDL (identity keys, deferred FK checks, +-- @DOUBLE PRECISION@). Dataset-specific fit tables are a later phase. +module Algorithm.EqSat.Storage.Schema + ( schemaSQL + , createSchema + ) where + +import Data.Text (Text) + +import Algorithm.EqSat.Storage.Backend (SqlBackend(..)) + +schemaSQL :: [Text] +schemaSQL = + [ "CREATE TABLE IF NOT EXISTS meta (" + <> " key TEXT PRIMARY KEY," + <> " value TEXT NOT NULL)" + , "CREATE TABLE IF NOT EXISTS enode (" + <> " key TEXT PRIMARY KEY," + <> " op TEXT NOT NULL," + <> " op_detail TEXT," + <> " a INTEGER," + <> " b INTEGER," + <> " x REAL)" + , "CREATE TABLE IF NOT EXISTS enode_child (" + <> " enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE," + <> " child_eid INTEGER NOT NULL," + <> " cnt INTEGER NOT NULL DEFAULT 1," + <> " PRIMARY KEY (enode_key, child_eid))" + , "CREATE TABLE IF NOT EXISTS eclass (" + <> " eid INTEGER PRIMARY KEY," + <> " canonical INTEGER NOT NULL," + <> " height INTEGER NOT NULL DEFAULT 0)" + , "CREATE TABLE IF NOT EXISTS eclass_node (" + <> " eid INTEGER NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE," + <> " PRIMARY KEY (eid, enode_key))" + , "CREATE TABLE IF NOT EXISTS parent (" + <> " child_eid INTEGER NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " parent_eid INTEGER NOT NULL," + <> " parent_enode_key TEXT NOT NULL REFERENCES enode(key) ON DELETE CASCADE," + <> " PRIMARY KEY (child_eid, parent_eid, parent_enode_key))" + , "CREATE TABLE IF NOT EXISTS cstore_page (" + <> " key TEXT PRIMARY KEY," + <> " blob TEXT NOT NULL)" + , "CREATE TABLE IF NOT EXISTS frontier (" + <> " eid INTEGER PRIMARY KEY REFERENCES eclass(eid) ON DELETE CASCADE," + <> " updated_at TEXT)" + , "CREATE TABLE IF NOT EXISTS dataset (" + <> " id INTEGER PRIMARY KEY," + <> " name TEXT NOT NULL UNIQUE," + <> " created TEXT)" + , "CREATE TABLE IF NOT EXISTS dataset_fit (" + <> " dataset_id INTEGER NOT NULL REFERENCES dataset(id) ON DELETE CASCADE," + <> " eid INTEGER NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " fitness REAL," + <> " dl REAL," + <> " theta TEXT," + <> " size INTEGER NOT NULL DEFAULT 0," + <> " evaluated INTEGER NOT NULL DEFAULT 0," + <> " fitted INTEGER NOT NULL DEFAULT 0," + <> " stale INTEGER NOT NULL DEFAULT 0," + <> " updated_at TEXT," + <> " PRIMARY KEY (dataset_id, eid))" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_fitness ON dataset_fit(fitness)" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_size ON dataset_fit(size)" + , "CREATE INDEX IF NOT EXISTS idx_dsfit_dl ON dataset_fit(dl)" + , "CREATE TABLE IF NOT EXISTS expression_index (" + <> " expression_key TEXT PRIMARY KEY," + <> " eclass INTEGER NOT NULL REFERENCES eclass(eid) ON DELETE CASCADE," + <> " dataset_id INTEGER REFERENCES dataset(id) ON DELETE CASCADE," + <> " first_seen TEXT)" + ] + +-- | Create (or ensure) the schema on the given backend. +createSchema :: SqlBackend db => db -> IO () +createSchema = createSchemaDb diff --git a/src/Algorithm/EqSat/Storage/Stream.hs b/src/Algorithm/EqSat/Storage/Stream.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Stream.hs @@ -0,0 +1,93 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Proof-of-concept for an O(1)-memory matcher: instead of building the +-- in-RAM pattern trie (@_patDB@, O(nodes)), stream the @enode@/@eclass_node@ +-- tables through a SQLite cursor and match each node as it is read. Peak memory +-- is O(budget) (the bounded result set), independent of graph size. +-- +-- This validates the mechanism behind P2.3 (the SQL/streaming matcher rewrite) +-- before committing to replacing the trie. It is SQLite-specific (uses +-- 'Database.SQLite3' directly, since 'SqlBackend' returns whole grids). +module Algorithm.EqSat.Storage.Stream + ( streamByOpCount + , streamMatchNAry + , streamRootsByOp + ) where + +import Database.SQLite3 + ( Database, SQLData(..), StepResult(..) + , bind, columns, step, withStatement ) +import qualified Data.Text as T +import qualified Data.IntSet as IntSet + +import Algorithm.EqSat.Egraph (EClassId, ENode(..)) +import Algorithm.EqSat.Storage.Types (parseEnodeKey) +import Data.Int (Int64) + +-- | Stream the @enode@ table by @op_detail@ and count rows without accumulating +-- them (the O(1)-memory baseline for a streaming matcher). +streamByOpCount :: Database -> T.Text -> IO Int +streamByOpCount db op = withStatement db + "SELECT e.key FROM enode e WHERE e.op_detail = ?" $ \stmt -> do + bind stmt [SQLText op] + go stmt 0 + where + go stmt n = do + r <- step stmt + case r of + Done -> pure n + Row -> go stmt (n + 1) + +-- | Stream the @eclass_node JOIN enode@ for a given operator, reconstruct each node +-- from its content key, and collect at most @budget@ e-class ids that actually +-- contain a node of that operator. Memory is O(@budget@), not O(nodes). +streamMatchNAry :: Database -> T.Text -> Int -> IO [EClassId] +streamMatchNAry db opBudget budget = withStatement db + "SELECT n.eid, n.enode_key FROM eclass_node n \ + \JOIN enode e ON e.key = n.enode_key WHERE e.op_detail = ?" $ \stmt -> do + bind stmt [SQLText opBudget] + go stmt budget [] + where + go stmt budgetLeft acc + | budgetLeft <= 0 = pure (reverse acc) + | otherwise = do + r <- step stmt + case r of + Done -> pure (reverse acc) + Row -> do + cols <- columns stmt + let eid = case cols of (SQLInteger i : _) -> fromIntegral i; _ -> 0 + key = case cols of (_ : SQLText k : _) -> T.unpack k; _ -> "" + ok = case parseEnodeKey key of + Just (ENAry _ _) -> True + Just _ -> True + Nothing -> False + if ok + then go stmt (budgetLeft - 1) (eid : acc) + else go stmt budgetLeft acc + +-- | Stream the distinct e-class ids that contain a node of a given @op_detail@ +-- through a SQLite cursor, stopping after @budget@ non-excluded rows. This is +-- the O(1)-memory candidate-root source for the streaming n-ary matcher: it +-- never materializes the whole (operator -> root set) index in RAM. Memory is +-- O(@budget@ + size of @exclude@). +streamRootsByOp :: Database -> T.Text -> Int -> [EClassId] -> IO [EClassId] +streamRootsByOp db opDetail budget exclude = withStatement db + "SELECT DISTINCT n.eid FROM eclass_node n \ + \JOIN enode e ON e.key = n.enode_key WHERE e.op_detail = ?" $ \stmt -> do + bind stmt [SQLText opDetail] + let ex = IntSet.fromList exclude + go stmt ex budget [] + where + go stmt ex n acc + | n <= 0 = pure (reverse acc) + | otherwise = do + r <- step stmt + case r of + Done -> pure (reverse acc) + Row -> do + cols <- columns stmt + let eid = case cols of (SQLInteger i : _) -> fromIntegral i; _ -> 0 + if IntSet.member eid ex + then go stmt ex n acc + else go stmt ex (n - 1) (eid : acc) diff --git a/src/Algorithm/EqSat/Storage/Types.hs b/src/Algorithm/EqSat/Storage/Types.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Storage/Types.hs @@ -0,0 +1,128 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Persistent representation helpers: canonical e-node content keys and +-- theta serialization used by the SQL storage layer. +module Algorithm.EqSat.Storage.Types + ( enodeKey + , parseEnodeKey + , enodeChildren + , enodeOpTag + , enodeOpDetail + , opDetailOf + , serializeTheta + , parseTheta + ) where + +import Data.SRTree.Internal (Op(..), Function(..), SRTree(..)) +import Data.SRTree.Eval (Target) +import Algorithm.EqSat.Egraph (ENode(..), EClassId, NOp(..)) + +import qualified Data.IntMap as IntMap +import qualified Data.Vector.Unboxed as VU +import Data.List (intercalate) +import Data.Char (isSpace) +import Text.Read (readMaybe) + +-- | Canonical, stable, parseable serialization of an e-node. +-- +-- Children reference e-class ids: for EUni/EBin they appear inline, for the +-- ENAry multiset they appear as @class:multiplicity@ pairs (sorted by class). +-- This string doubles as the content-addressable primary key of @enode@. +enodeKey :: ENode -> String +enodeKey (EVar i) = "Var " <> show i +enodeKey (EParam i) = "Param " <> show i +enodeKey (EConst x) = "Const " <> show x +enodeKey (EUni f c) = "Uni " <> show f <> " " <> show c +enodeKey (EBin op l r) = "Bin " <> show op <> " " <> show l <> " " <> show r +enodeKey (ENAry op m) = "NAry " <> showNOp op <> " " <> intercalate " " [show c <> ":" <> show n | (c, n) <- IntMap.toAscList m] + +showNOp :: NOp -> String +showNOp EAdd = "EAdd" +showNOp EMul = "EMul" + +parseEnodeKey :: String -> Maybe ENode +parseEnodeKey = go . words + where + go ["Var", i] = EVar <$> readMaybe i + go ["Param", i] = EParam <$> readMaybe i + go ["Const", x] = EConst <$> readMaybe x + go ["Uni", f, c] = do + fc <- readMaybe f + cc <- readMaybe c + pure (EUni fc cc) + go ["Bin", op, l, r] = do + oc <- readMaybe op + lc <- readMaybe l + rc <- readMaybe r + pure (EBin oc lc rc) + go ("NAry" : op : rest) = do + oc <- parseNOp op + ch <- mapM parseChild rest + pure (ENAry oc (IntMap.fromList ch)) + go _ = Nothing + + parseChild w = case break (== ':') w of + (c, ':' : n) -> (,) <$> readMaybe c <*> readMaybe n + _ -> Nothing + +parseNOp :: String -> Maybe NOp +parseNOp "EAdd" = Just EAdd +parseNOp "EMul" = Just EMul +parseNOp _ = Nothing + +-- | All e-class ids referenced as children by a node. +enodeChildren :: ENode -> [EClassId] +enodeChildren (EUni _ c) = [c] +enodeChildren (EBin _ l r) = [l, r] +enodeChildren (ENAry _ m) = IntMap.keys m +enodeChildren _ = [] + +-- | Coarse operator tag used as the @op@ column: Var | Param | Const | Uni | Bin | NAry. +enodeOpTag :: ENode -> String +enodeOpTag (EVar _) = "Var" +enodeOpTag (EParam _) = "Param" +enodeOpTag (EConst _) = "Const" +enodeOpTag (EUni _ _) = "Uni" +enodeOpTag (EBin _ _ _) = "Bin" +enodeOpTag (ENAry _ _) = "NAry" + +-- | Specific operator detail (EAdd/EMul/Add/Sub/.../LogAbs/...) for pattern counting. +enodeOpDetail :: ENode -> String +enodeOpDetail (EVar _) = "Var" +enodeOpDetail (EParam _) = "Param" +enodeOpDetail (EConst _) = "Const" +enodeOpDetail (EUni f _) = show f +enodeOpDetail (EBin op _ _) = show op +enodeOpDetail (ENAry op _) = showNOp op + +-- | The @op_detail@ column value matching an operator-key shape (@SRTree ()@, +-- as produced by 'Algorithm.EqSat.Egraph.eOpKey' / pattern @opOf@). N-ary +-- Add/Mul patterns address the flattened ENAry nodes, whose @op_detail@ is +-- @EAdd@/@EMul@, so the binary shapes @Bin Add@/@Bin Mul@ map there; other +-- operators map to their own detail. +opDetailOf :: SRTree () -> String +opDetailOf (Var _) = "Var" +opDetailOf (Param _) = "Param" +opDetailOf (Const _) = "Const" +opDetailOf (Y _) = "Var" +opDetailOf (Uni f _) = show f +opDetailOf (Bin Add _ _) = showNOp EAdd +opDetailOf (Bin Mul _ _) = showNOp EMul +opDetailOf (Bin op _ _) = show op + +-- | Flatten a list of target vectors: vectors joined by @|@, elements by @,@. +serializeTheta :: [Target] -> String +serializeTheta ts = intercalate "|" [ intercalate "," (map show (VU.toList t)) | t <- ts ] + +-- | Inverse of 'serializeTheta'. +parseTheta :: String -> [Target] +parseTheta s = [ VU.fromList (map readDouble (splitOn ',' v)) | v <- splitOn '|' s, not (null (filter (not . isSpace) v)) ] + where + readDouble x = case readMaybe x of + Just d -> d + Nothing -> 0 + +splitOn :: Char -> String -> [String] +splitOn c s = case break (== c) s of + (a, [] ) -> if null a then [] else [a] + (a, _ : b) -> a : splitOn c b diff --git a/srtree-db.cabal b/srtree-db.cabal new file mode 100644 --- /dev/null +++ b/srtree-db.cabal @@ -0,0 +1,88 @@ +cabal-version: 2.4 +name: srtree-db +version: 0.1.0.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 + rows (backed by SQLite and PostgreSQL) plus SQL queries + (topN, pareto, distribution counts, patterns). + Dataset-specific fit tables are a later phase. +license: BSD-3-Clause +license-file: LICENSE +author: Fabricio Olivetti de França +maintainer: fabricio.olivetti@gmail.com +category: Math, Data +build-type: Simple +homepage: https://github.com/folivetti/srtree-db#readme +bug-reports: https://github.com/folivetti/srtree-db/issues +extra-source-files: + ChangeLog.md + +source-repository head + type: git + location: https://github.com/folivetti/srtree-db + +library + hs-source-dirs: src + exposed-modules: + Algorithm.EqSat.Storage.Types + Algorithm.EqSat.Storage.Backend + Algorithm.EqSat.Storage.ClassStore + Algorithm.EqSat.Storage.Import + Algorithm.EqSat.Storage.Stream + Algorithm.EqSat.Storage.Schema + Algorithm.EqSat.Storage.SQLite + Algorithm.EqSat.Storage.Postgres + Algorithm.EqSat.Storage.Query + build-depends: + base >=4.14 && <5 + , bytestring >=0.10 && <0.13 + , binary >=0.8 && <0.9 + , containers >=0.6 && <0.9 + , unordered-containers >=0.2 && <0.3 + , text >=1.2 && <2.2 + , direct-sqlite >=2.3 && <2.4 + , postgresql-libpq >=0.10 && <0.12 + , srtree >=3.0 && <3.1 + , vector >=0.12 && <0.14 + , mtl >=2.2 && <2.4 + default-language: Haskell2010 + +benchmark srtree-db-bench + type: exitcode-stdio-1.0 + main-is: Bench.hs + hs-source-dirs: bench + ghc-options: -O2 + build-depends: + base >=4.14 && <5 + , bytestring >=0.10 && <0.13 + , containers >=0.6 && <0.9 + , text >=1.2 && <2.2 + , time >=1.10 && <1.15 + , directory >=1.3 && <1.4 + , direct-sqlite >=2.3 && <2.4 + , postgresql-libpq >=0.10 && <0.12 + , srtree >=3.0 && <3.1 + , srtree-db >=0.1 && <0.2 + default-language: Haskell2010 + +test-suite srtree-db-test + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test + build-depends: + base >=4.14 && <5 + , HUnit >=1.6 && <1.7 + , bytestring >=0.10 && <0.13 + , binary >=0.8 && <0.9 + , containers >=0.6 && <0.9 + , directory >=1.3 && <1.4 + , unordered-containers >=0.2 && <0.3 + , text >=1.2 && <2.2 + , direct-sqlite >=2.3 && <2.4 + , postgresql-libpq >=0.10 && <0.12 + , srtree >=3.0 && <3.1 + , srtree-db >=0.1 && <0.2 + , vector >=0.12 && <0.14 + , mtl >=2.2 && <2.4 + default-language: Haskell2010 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,597 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Main where + +import Test.HUnit +import Control.Monad (forM_) +import Control.Exception (catch, SomeException) +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Identity (runIdentity, Identity) +import Control.Monad.State.Strict (runStateT, execStateT, evalStateT, get) +import qualified Data.IntMap as IntMap +import qualified Data.HashSet as Set +import qualified Data.ByteString as BS +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Vector.Unboxed as VU +import Data.Text (Text) +import Data.SRTree.Eval (Target) +import System.Environment (lookupEnv) +import System.Directory (removeFile) +import Data.Maybe (isJust, isNothing) + +import Database.SQLite3 ( close, open ) +import Database.PostgreSQL.LibPQ ( Connection, connectdb, finish ) + +import Data.SRTree +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Simplify (rewrites) +import Algorithm.EqSat.Build (fromTree) +import Algorithm.EqSat.Info (insertFitness) +import Algorithm.EqSat.DB (Pattern(..), match) +import Algorithm.EqSat.Queries (findRootClasses) +import Algorithm.EqSat.Storage.Backend ( SqlValue(..), SqlBackend(..), sqlToInt ) +import Algorithm.EqSat.Storage.ClassStore +import Algorithm.EqSat.Storage.SQLite +import Algorithm.EqSat.Storage.Query (getOrCreateDataset) +import Algorithm.EqSat.Storage.Backend (SqlBackend) +import Algorithm.EqSat.Storage.Postgres () +import qualified Algorithm.EqSat.Storage.Query as Q + +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 + +runIn :: EGraph -> EGraphST Identity a -> (a, EGraph) +runIn g m = runIdentity $ runStateT m g + +evalIn :: EGraph -> EGraphST Identity a -> a +evalIn g m = runIdentity $ evalStateT m g + +-- | Run a graph action in an IO monad so that a paged store (when the graph +-- carries one) is exercised write-through. +runIOIn :: EGraph -> EGraphST IO a -> IO (a, EGraph) +runIOIn g m = runStateT m g + +-- | 'saveGraph' now needs a dataset id; give the tests a fixed dataset. +saveGraphTest :: SqlBackend db => db -> EGraph -> IO (Either String ()) +saveGraphTest db eg = do + dsid <- getOrCreateDataset db "test" + saveGraph db dsid eg + +-- | x0..x3; add = x0+x1 (fit 0.9), p2 = (x0+x1)*x2 (fit 0.7), +-- p3 = (x0+x1)*x3 (fit 0.5). Returns (eg, add, p2, p3). +buildGraph :: IO (EGraph, EClassId, EClassId, EClassId) +buildGraph = pure (eg, eidAdd, eidP2, eidP3) + where + ((eidAdd, eidP2, eidP3), eg) = runIn emptyGraph go + go = do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + eidAdd <- fromTree myCost (var 0 + var 1) + eidP2 <- fromTree myCost ((var 0 + var 1) * var 2) + eidP3 <- fromTree myCost ((var 0 + var 1) * var 3) + insertFitness eidAdd 0.9 [] + insertFitness eidP2 0.7 [] + insertFitness eidP3 0.5 [VU.fromList [1.0, 2.0]] + pure (eidAdd, eidP2, eidP3) + +numEvaluated :: EGraph -> Int +numEvaluated eg = length + [ () | (_, ec) <- IntMap.toList (_eClass eg), _fitness (_info ec) /= Nothing ] + +-- --------------------------------------------------------------------------- +-- driver-generic tests (run once against SQLite, once against PostgreSQL) + +testSaveLoadRT :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testSaveLoadRT openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, eidP3) <- buildGraph + _ <- saveGraphTest db eg + m <- loadGraph db + case m of + Left err -> assertFailure ("loadGraph failed: " <> err) + Right eg' -> do + assertEqual "class count preserved" (IntMap.size (_eClass eg)) (IntMap.size (_eClass eg')) + assertEqual "fitness preserved" (Just 0.9) (evalIn eg' (getFitness eidAdd)) + -- theta round-trips on the class that had params + let th = evalIn eg' (getTheta eidP3) + assertEqual "theta preserved" 1 (length th) + -- (x0+x1)*(x2|x3) must still be matchable after import + let prod = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'A') (VarPat 'B'))) (VarPat 'C')) + assertBool "patterns queryable after load" + (length (evalIn eg' (match prod)) >= 2) + -- inserting the already-loaded (x0+x1) dedups (no new class) + let (eid2, g2) = runIn eg' (fromTree myCost (var 0 + var 1)) + assertEqual "post-load dedup" eidAdd eid2 + assertEqual "post-load no growth" (IntMap.size (_eClass eg)) (IntMap.size (_eClass g2)) + closeDb db + +testQueries :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testQueries openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, eidP2, _) <- buildGraph + _ <- saveGraphTest db eg + tn2 <- Q.topN db 1 2 + case tn2 of + (x : _) -> assertEqual "topN[0]" (eidAdd, 0.9) x + [] -> assertFailure "topN returned no rows" + assertEqual "topN length" 2 (length tn2) + tn1 <- Q.topN db 1 1 + assertEqual "topN(1)" [(eidAdd, 0.9)] tn1 + dc <- Q.distributionCounts db 1 100 + assertEqual "distribution totals evaluated classes" (numEvaluated eg) (sum (map snd dc)) + cAdd <- Q.countPattern db "EAdd" + assertEqual "EAdd classes" 1 cAdd + cMul <- Q.countPattern db "EMul" + assertEqual "EMul classes" 2 cMul + -- pareto over (max fitness, min dl): give add (fit 0.9, dl 1.0) which + -- dominates p2 (0.7, 2.0) and p3 (0.5, 3.0), so only add remains. + execDb db ("UPDATE dataset_fit SET dl = 1.0 WHERE eid = " <> T.pack (show eidAdd)) + execDb db ("UPDATE dataset_fit SET dl = 2.0 WHERE eid = " <> T.pack (show eidP2)) + execDb db ("UPDATE dataset_fit SET dl = 3.0 WHERE eid NOT IN (" <> T.pack (show eidAdd) <> "," <> T.pack (show eidP2) <> ")") + p <- Q.pareto db 1 + assertEqual "pareto keeps only add" [(eidAdd, 0.9, 1.0)] p + closeDb db + +testSync :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testSync openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, _) <- buildGraph + _ <- saveGraphTest db eg + + -- refit in memory, push to DB + let egNew = runIdentity $ execStateT (insertFitness eidAdd 0.99 []) eg + pushFit db 1 egNew + Right eg3 <- refreshFitness db 1 eg + assertEqual "pushed fitness read back" (Just 0.99) (evalIn eg3 (getFitness eidAdd)) + assertEqual "graph class count intact" (IntMap.size (_eClass eg)) (IntMap.size (_eClass eg3)) + + -- edit fitness straight in the DB, refresh pulls it in; refreshFitness and + -- loadGraph both read dataset_fit now (the legacy fit table is gone). + execDb db ("UPDATE dataset_fit SET fitness = 0.55 WHERE eid = " <> T.pack (show eidAdd)) + Right eg4 <- refreshFitness db 1 eg + assertEqual "db edit pulled in" (Just 0.55) (evalIn eg4 (getFitness eidAdd)) + + -- a stored graph loads with the current DB fitness + Right eg5 <- loadGraph db + assertEqual "load uses current DB fitness" (Just 0.55) (evalIn eg5 (getFitness eidAdd)) + closeDb db + +pgOpen :: String -> IO Connection +pgOpen dsn = connectdb (TE.encodeUtf8 (T.pack dsn)) + +-- --------------------------------------------------------------------------- +-- ClassStore tests (page round-trip, persistence, LRU bound, delete) + +mkPage :: Int -> BS.ByteString +mkPage i = BS.replicate (10 + i) (fromIntegral i) + +-- | The @parent@ table round-trips every reverse edge: the row count matches +-- the live graph's total @_parents@ entries, and a reload reconstructs the +-- identical per-class parent sets. +testParents :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testParents openDb closeDb = TestCase $ do + db <- openDb + (eg, _, _, _) <- buildGraph + _ <- saveGraphTest db eg + let expected = sum + [ Set.size (_parents ec) + | (_, ec) <- IntMap.toList (_eClass eg) ] + rows <- query db "SELECT COUNT(*) FROM parent" [] + case rows of + [[SqlInteger n]] -> assertEqual "parent row count" expected (fromIntegral n) + [[SqlText t]] -> assertEqual "parent row count" expected + (read (T.unpack t) :: Int) + _ -> assertFailure "parent count: unexpected row shape" + Right eg' <- loadGraph db + let parentsMap g = IntMap.map _parents (_eClass g) + assertEqual "parents preserved per class" (parentsMap eg) (parentsMap eg') + closeDb db + +testStoreRoundtrip :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testStoreRoundtrip openDb closeDb = TestCase $ do + db <- openDb + execDb db "DROP TABLE IF EXISTS cstore_test" + ps <- newPageStore db "cstore_test" 0 100 + forM_ [0 .. 4] $ \i -> writePage ps i (mkPage i) + forM_ [0 .. 4] $ \i -> do + mp <- readPage ps i + assertEqual ("page " <> show i) (Just (mkPage i)) mp + assertEqual "pending before flush" 5 =<< pendingCount ps + assertEqual "writeback count" 5 =<< writeback ps + assertEqual "pending after flush" 0 =<< pendingCount ps + closeDb db + +testStorePersist :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testStorePersist openDb closeDb = TestCase $ do + db <- openDb + execDb db "DROP TABLE IF EXISTS cstore_test" + ps1 <- newPageStore db "cstore_test" 0 100 + writePage ps1 42 (mkPage 42) + _ <- writeback ps1 + closeDb db + db2 <- openDb + ps2 <- newPageStore db2 "cstore_test" 0 100 + mp <- readPage ps2 42 + assertEqual "persisted page" (Just (mkPage 42)) mp + closeDb db2 + +testStoreLRU :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testStoreLRU openDb closeDb = TestCase $ do + db <- openDb + execDb db "DROP TABLE IF EXISTS cstore_test" + ps <- newPageStore db "cstore_test" 2 100 + forM_ [0 .. 9] $ \i -> writePage ps i (mkPage i) + assertEqual "resident bounded by cap" 2 =<< residentCount ps + _ <- writeback ps + forM_ [0 .. 9] $ \i -> do + mp <- readPage ps i + assertEqual ("evicted reload " <> show i) (Just (mkPage i)) mp + deletePage ps 3 + mp <- readPage ps 3 + assertEqual "deleted page gone" Nothing mp + closeDb db + +runStoreSuite :: SqlBackend db => String -> (IO db, db -> IO ()) -> [Test] +runStoreSuite tag (openDb, closeDb) = + [ TestLabel (tag <> " store-roundtrip") (testStoreRoundtrip openDb closeDb) + , TestLabel (tag <> " store-persist") (testStorePersist openDb closeDb) + , TestLabel (tag <> " store-lru") (testStoreLRU openDb closeDb) + ] + +-- | The e-class page store round-trips: saveGraph writes one page per class, +-- loadGraph installs the paged store handle, class count/fitness/parents are +-- preserved, and post-load dedup does not grow the graph. +testPagedRoundtrip :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testPagedRoundtrip openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, eidP3) <- buildGraph + _ <- saveGraphTest db eg + pc <- query db "SELECT COUNT(*) FROM cstore_page" [] + let expected = IntMap.size (_eClass eg) + case pc of + [[SqlInteger n]] -> assertEqual "page count" expected (fromIntegral n) + [[SqlText t]] -> assertEqual "page count" expected (read (T.unpack t)) + _ -> assertFailure "page count: unexpected row shape" + m <- loadGraph db + case m of + Left err -> assertFailure ("loadGraph failed: " <> err) + Right eg' -> do + assertBool "paged store installed" (isJust (_classStore eg')) + assertEqual "class count preserved" (IntMap.size (_eClass eg)) (IntMap.size (_eClass eg')) + assertEqual "fitness (paged)" (Just 0.9) (evalIn eg' (getFitness eidAdd)) + assertEqual "theta (paged)" 1 (length (evalIn eg' (getTheta eidP3))) + -- inserting an already-loaded node dedups + let (eid2, g2) = runIn eg' (fromTree myCost (var 0 + var 1)) + assertEqual "post-load dedup" eidAdd eid2 + assertEqual "post-load no growth" (IntMap.size (_eClass eg)) (IntMap.size (_eClass g2)) + -- flushing a paged graph is safe + flushStore eg' + closeDb db + +-- | Mutations through the paged store persist: add a class + fitness, save +-- again, reload, and both the new and old classes are present. +testPagedMutatePersist :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testPagedMutatePersist openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, _) <- buildGraph + _ <- saveGraphTest db eg + Right eg0 <- loadGraph db + (eidNew, eg1) <- runIOIn eg0 $ do + cid <- fromTree myCost (var 0 + var 2) + insertFitness cid 0.33 [] + st <- get + liftIO (flushStore st) + pure cid + -- (x0+x2) was not present, so a brand new class was created + assertBool "new class is distinct" (eidNew /= eidAdd) + _ <- saveGraphTest db eg1 + Right eg2 <- loadGraph db + assertEqual "new class fitness persisted" (Just 0.33) (evalIn eg2 (getFitness eidNew)) + assertEqual "old class intact" (Just 0.9) (evalIn eg2 (getFitness eidAdd)) + assertEqual "paged store reinstalled" True (isJust (_classStore eg2)) + closeDb db + +-- | A database without pages (page table emptied) loads through the fully +-- relational path, with no paged store handle. +testPagedFallback :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testPagedFallback openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, _) <- buildGraph + _ <- saveGraphTest db eg + execDb db "DELETE FROM cstore_page" + m <- loadGraph db + case m of + Left err -> assertFailure ("fallback load failed: " <> err) + Right eg' -> do + assertBool "fallback: no paged store" (isNothing (_classStore eg')) + assertEqual "fallback class count" (IntMap.size (_eClass eg)) (IntMap.size (_eClass eg')) + assertEqual "fallback fitness" (Just 0.9) (evalIn eg' (getFitness eidAdd)) + closeDb db + +-- | The out-of-core path: 'loadGraphLazy' leaves the resident e-class map +-- empty, installs the paged store, and streams every whole-graph read through +-- the store (bounds memory). Mutations persist through the store. +testLazyLoad :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testLazyLoad openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, eidP2, _) <- buildGraph + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg' -> do + assertBool "lazy: paged store installed" (isJust (_classStore eg')) + -- resident cache starts empty (nothing materialized) + assertEqual "lazy: resident cache empty" 0 (IntMap.size (_eClass eg')) + -- whole-graph enumeration streams from the store + (classes, g1) <- runIOIn eg' allClasses + assertEqual "lazy: allClasses accepts-every-class" + (IntMap.size (_eClass eg)) (length classes) + -- point reads fall back to the store + (f, _) <- runIOIn g1 (getFitness eidAdd) + assertEqual "lazy: getFitness from store" (Just 0.9) f + -- structural pattern matching works on the seeded patDB/range DBs + let pat = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'A') (VarPat 'B'))) (VarPat 'C')) + (ms, _) <- runIOIn g1 (match pat) + assertBool "lazy: patterns queryable" (length ms >= 2) + -- root enumeration streams from the store + (roots, _) <- runIOIn g1 findRootClasses + assertEqual "lazy: two root classes" 2 (length roots) + -- mutate through the lazy graph, then persist the whole (paged) graph + (eidNew, eg2) <- runIOIn g1 $ do + cid <- fromTree myCost (var 3 + var 0) -- a fresh class, not in the DB + insertFitness cid 0.11 [] + st <- get + liftIO (flushStore st) + pure cid + _ <- saveGraphTest db eg2 + Right eg3 <- loadGraph db + assertEqual "lazy mutate: new class persisted" (Just 0.11) (evalIn eg3 (getFitness eidNew)) + assertEqual "lazy mutate: old class intact" (Just 0.9) (evalIn eg3 (getFitness eidAdd)) + closeDb db + +-- | A lazy (resident-empty) graph can still be rewritten: 'runEqSat' streams +-- every class through the paged store while evaluating the const-valued +-- preconditions (e.g. 'isNotZero') against stored class info. After a pass with +-- the 'rewrites' ruleset, the store-aware conditions fired and the constant +-- rewrites (x-x->0, 1**x->1, x/x->1) landed. +testLazyRewrite :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testLazyRewrite openDb closeDb = TestCase $ do + db <- openDb + let eg = snd $ runIn emptyGraph $ do + _ <- fromTree myCost (var 0 - var 0) + _ <- fromTree myCost (1 ** var 1) + _ <- fromTree myCost (var 0 / var 0) + pure () + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg' -> do + assertEqual "lazy: resident cache empty" 0 (IntMap.size (_eClass eg')) + -- one eqsat pass; conditions are evaluated against the paged store + (_, g1) <- runIOIn eg' (runEqSat myCost rewrites 10) + -- x - x rewrote to Const 0 + (m0, _) <- runIOIn g1 (match (Fixed (Const 0))) + assertBool "lazy rewrite: x-x -> 0" (not (null m0)) + -- 1 ** x and x / x (guarded by isNotZero) rewrote to Const 1 + (m1, _) <- runIOIn g1 (match (Fixed (Const 1))) + assertBool "lazy rewrite: produced Const 1" (length m1 >= 1) + flushStore g1 + closeDb db + +-- | Out-of-core round trip through eqsat. Saturate a paged graph (which +-- *merges* classes and writes the merged bodies + canonical mappings through +-- the write-through store), save it (paged save writes meta only), then reload +-- and check the saturated graph -- including the merge and its canonical +-- mapping -- is reconstructed from the live write-through tables + pages. +-- +-- This is the correctness guard for P2.2: with a bounded resident cache the +-- page store is the only authoritative copy of merged class bodies and the +-- `eclass.canonical` write-through is the only record of merges, so a reload +-- must see them. +testPagedEqSatReload :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testPagedEqSatReload openDb closeDb = TestCase $ do + db <- openDb + -- x0**2 and x0*x0 are distinct classes; the default rewrite x*x -> x**2 + -- merges them during eqsat. + let ((eidP, eidXx, eidAdd), eg) = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + eidP <- fromTree myCost (var 0 ** 2) + eidXx <- fromTree myCost (var 0 * var 0) + eidAdd <- fromTree myCost (var 0 + var 1) + insertFitness eidAdd 0.9 [] + pure (eidP, eidXx, eidAdd) + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg0 -> do + -- saturate; x0*x0 should merge into x0**2 + (_, g1) <- runIOIn eg0 (runEqSat myCost rewrites 30) + (cXx, _) <- runIOIn g1 (canonical eidXx) + (cP, _) <- runIOIn g1 (canonical eidP) + assertEqual "eqsat merged x0*x0 into x0**2" cP cXx + (ksAfter, _) <- runIOIn g1 allKeys + let nAfter = length ksAfter + -- persist: paged save writes meta only; write-through keeps pages and + -- the canonical/node tables live. + flushStore g1 + -- enode_child is populated for the nodes created during eqsat + ecRows <- query db "SELECT COUNT(*) FROM enode_child" [] + assertBool "enode_child populated during eqsat" + (case ecRows of [[v]] -> sqlToInt v > 0; _ -> False) + _ <- saveGraphTest db g1 + -- reload a fresh lazy graph + obj2 <- loadGraphLazy db 1 + case obj2 of + Left err2 -> assertFailure ("reload failed: " <> err2) + Right eg2 -> do + assertEqual "reload: resident cache empty" 0 (IntMap.size (_eClass eg2)) + -- same number of classes: the merge persisted (the absorbed class's + -- page was deleted, its nodes/canonical remain queryable) + (ksReload, _) <- runIOIn eg2 allKeys + assertEqual "class count persists across reload" nAfter (length ksReload) + -- the merged canonical mapping persists (resolves with no CANON_MISSING) + (cXx2, _) <- runIOIn eg2 (canonical eidXx) + (cP2, _) <- runIOIn eg2 (canonical eidP) + assertEqual "merged canonical persists" cXx2 cP2 + -- the surviving (merged) class still matches its expression shape + (m, _) <- runIOIn eg2 (match (Fixed (Bin Power (VarPat 'A') (VarPat 'B')))) + assertBool "x**y matchable after reload" (not (null m)) + -- fitness persisted via the page blob / dataset_fit + (f, _) <- runIOIn eg2 (getFitness eidAdd) + assertEqual "fitness persists" (Just 0.9) f + closeDb db + +-- | The paged (streamed) 'pushFit' path: change fitness in a lazy graph, flush, +-- push to @dataset_fit@, and read it back via 'refreshFitness'. +testPagedPushFit :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testPagedPushFit openDb closeDb = TestCase $ do + db <- openDb + (eg, eidAdd, _, _) <- buildGraph + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg0 -> do + (_, g1) <- runIOIn eg0 (insertFitness eidAdd 0.99 []) + flushStore g1 + pushFit db 1 g1 -- streamed (O(1) memory) push + Right eg2 <- refreshFitness db 1 eg + assertEqual "paged push fitness read back" (Just 0.99) (evalIn eg2 (getFitness eidAdd)) + closeDb db + +-- | Frontier re-saturation: seed the @frontier@ table with only one of two +-- mergeable pairs, run a frontier pass, and check that (a) the frontier class +-- is re-saturated (its merge happens) while a mergeable pair outside the +-- frontier is left untouched, and (b) the frontier is cleared afterwards. This +-- is the "re-saturate only recently-changed classes, avoid redo" behaviour; the +-- pure in-memory eggp loop and a full dbEqSat are unaffected. +testFrontierReload :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testFrontierReload openDb closeDb = TestCase $ do + db <- openDb + let ((eidA, eidB, eidC, eidD), eg) = runIn emptyGraph $ do + eidA <- fromTree myCost (var 0 * var 0) -- x0*x0 (mergeable, NOT in frontier) + eidB <- fromTree myCost (var 0 ** 2) -- x0**2 + eidC <- fromTree myCost (var 1 * var 1) -- x1*x1 (mergeable, IN frontier) + eidD <- fromTree myCost (var 1 ** 2) -- x1**2 + pure (eidA, eidB, eidC, eidD) + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg0 -> case _classStore eg0 of + Nothing -> assertFailure "frontier: expected a paged store" + Just h -> do + runDb db "INSERT INTO frontier (eid) VALUES (?)" [SqlInteger (fromIntegral eidC)] + (_, g1) <- runIOIn eg0 $ do + liftIO (cpsBeginFrontier h) + _ <- runEqSat myCost rewrites 20 + liftIO (cpsEndFrontier h) + pure () + -- the frontier class was re-saturated: x1*x1 merged into x1**2 + (cC, _) <- runIOIn g1 (canonical eidC) + (cD, _) <- runIOIn g1 (canonical eidD) + assertEqual "frontier: x1*x1 merged into x1**2" cD cC + -- a mergeable pair OUTSIDE the frontier was left untouched + (cA, _) <- runIOIn g1 (canonical eidA) + (cB, _) <- runIOIn g1 (canonical eidB) + assertBool "frontier: x0*x0 NOT merged (outside frontier)" (cA /= cB) + -- the frontier was cleared after the pass + fr <- loadFrontierRows db + assertBool "frontier cleared after pass" (null fr) + closeDb db + +-- | Equivalence of the pure in-memory eqsat and the paged DB eqsat: running the +-- same expressions through both must produce the same merge structure (which +-- classes collapse into one). e-class ids differ between the two graphs, so we +-- compare the *relative* structure -- A/B merge and are distinct from C/D -- in +-- each mode independently. This guards the "both options" guarantee: the DB +-- work didn't change in-memory saturation quality, and vice-versa. +testEquivInMemDB :: SqlBackend db => IO db -> (db -> IO ()) -> Test +testEquivInMemDB openDb closeDb = TestCase $ do + db <- openDb + let ((eidA, eidB, eidC, eidD), eg) = runIn emptyGraph $ do + eidA <- fromTree myCost (var 0 * var 0) -- mergeable pair #1 + eidB <- fromTree myCost (var 0 ** 2) + eidC <- fromTree myCost (var 1 * var 1) -- mergeable pair #2 + eidD <- fromTree myCost (var 1 ** 2) + pure (eidA, eidB, eidC, eidD) + -- --- in-memory path (pure Identity graph) --- + let (_, egMem) = runIn eg (runEqSat myCost rewrites 20) + (mA, _) = runIn egMem (canonical eidA) + (mB, _) = runIn egMem (canonical eidB) + (mC, _) = runIn egMem (canonical eidC) + (mD, _) = runIn egMem (canonical eidD) + assertEqual "in-mem: pair1 merged" mA mB + assertEqual "in-mem: pair2 merged" mC mD + assertBool "in-mem: pairs distinct" (mA /= mC) + -- --- DB/paged path --- + _ <- saveGraphTest db eg + obj <- loadGraphLazy db 1 + case obj of + Left err -> assertFailure ("loadGraphLazy failed: " <> err) + Right eg0 -> do + (_, egDB) <- runIOIn eg0 (runEqSat myCost rewrites 20) + (dA, _) <- runIOIn egDB (canonical eidA) + (dB, _) <- runIOIn egDB (canonical eidB) + (dC, _) <- runIOIn egDB (canonical eidC) + (dD, _) <- runIOIn egDB (canonical eidD) + assertEqual "DB: pair1 merged" dA dB + assertEqual "DB: pair2 merged" dC dD + assertBool "DB: pairs distinct" (dA /= dC) + closeDb db + +runPagedSuite :: SqlBackend db => String -> (IO db, db -> IO ()) -> [Test] +runPagedSuite tag (openDb, closeDb) = + [ TestLabel (tag <> " paged-roundtrip") (testPagedRoundtrip openDb closeDb) + , TestLabel (tag <> " paged-mutate-persist") (testPagedMutatePersist openDb closeDb) + , TestLabel (tag <> " paged-fallback") (testPagedFallback openDb closeDb) + , TestLabel (tag <> " lazy-load") (testLazyLoad openDb closeDb) + , TestLabel (tag <> " lazy-rewrite") (testLazyRewrite openDb closeDb) + , TestLabel (tag <> " paged-push-fit") (testPagedPushFit openDb closeDb) + , TestLabel (tag <> " paged-eqsat-reload") (testPagedEqSatReload openDb closeDb) + , TestLabel (tag <> " frontier") (testFrontierReload openDb closeDb) + , TestLabel (tag <> " equiv-inmem-db") (testEquivInMemDB openDb closeDb) + ] + +runSuite :: SqlBackend db => String -> (IO db, db -> IO ()) -> [Test] +runSuite tag (openDb, closeDb) = + [ TestLabel (tag <> " save-load-roundtrip") (testSaveLoadRT openDb closeDb) + , TestLabel (tag <> " parents") (testParents openDb closeDb) + , TestLabel (tag <> " queries") (testQueries openDb closeDb) + , TestLabel (tag <> " sync") (testSync openDb closeDb) + ] + <> runStoreSuite tag (openDb, closeDb) + +main :: IO () +main = do + let sqlitePath = "/tmp/opencode/srtree-db-test.sqlite" + removeFile sqlitePath `catch` (\(_ :: SomeException) -> pure ()) + sqliteCounts <- runTestTT $ TestList + (runSuite "sqlite" (open (T.pack sqlitePath), close) + <> runPagedSuite "sqlite" (open (T.pack sqlitePath), close)) + mdsn <- lookupEnv "PGDSN" + case mdsn of + Nothing -> do + if failures sqliteCounts /= 0 || errors sqliteCounts /= 0 + then error "Some tests failed" + else do + putStrLn "PGDSN not set -- skipping PostgreSQL backend tests" + pure () + Just dsn -> do + pgCounts <- runTestTT $ TestList + (runSuite "postgresql" (pgOpen dsn, finish) + <> runPagedSuite "postgresql" (pgOpen dsn, finish)) + if failures sqliteCounts /= 0 || errors sqliteCounts /= 0 + || failures pgCounts /= 0 || errors pgCounts /= 0 + then error "Some tests failed" + else pure ()