agent: add build_relations_vector function to sqlite

Author: spaced4ndy

simplexmq.cabal

@@ -90,6 +90,7 @@ library
       Simplex.FileTransfer.Transport
       Simplex.FileTransfer.Types
       Simplex.FileTransfer.Util
+      Simplex.MemberRelations.MemberRelations
       Simplex.Messaging.Agent
       Simplex.Messaging.Agent.Client
       Simplex.Messaging.Agent.Env.SQLite

src/Simplex/MemberRelations/MemberRelations.hs

@@ -0,0 +1,83 @@
+{-# LANGUAGE LambdaCase #-}
+
+module Simplex.MemberRelations.MemberRelations
+  ( MemberRelation (..),
+    toRelationInt,
+    fromRelationInt,
+    getRelation,
+    getRelationsIndexes,
+    setRelation,
+    setRelations,
+  )
+where
+
+import Control.Monad
+import Data.Bits ((.&.), (.|.), complement)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import Data.ByteString.Internal (toForeignPtr, unsafeCreate)
+import Data.Int (Int64)
+import Data.Word (Word8)
+import Foreign.ForeignPtr (withForeignPtr)
+import Foreign.Marshal.Utils (copyBytes, fillBytes)
+import Foreign.Ptr (plusPtr)
+import Foreign.Storable (peekByteOff, pokeByteOff)
+
+data MemberRelation
+  = MRNew
+  | MRIntroduced
+  | MRIntroducedTo
+  | MRConnected
+  deriving (Eq, Show)
+
+toRelationInt :: MemberRelation -> Word8
+toRelationInt = \case
+  MRNew -> 0
+  MRIntroduced -> 1
+  MRIntroducedTo -> 2
+  MRConnected -> 3
+
+fromRelationInt :: Word8 -> MemberRelation
+fromRelationInt = \case
+  0 -> MRNew
+  1 -> MRIntroduced
+  2 -> MRIntroducedTo
+  3 -> MRConnected
+  _ -> MRNew
+
+-- | Get the relation status of a member at a given index from the relations vector.
+-- Returns 'MRNew' if the vector is not long enough (lazy initialization).
+getRelation :: Int64 -> ByteString -> MemberRelation
+getRelation i v
+  | i < 0 || fromIntegral i >= B.length v = MRNew
+  | otherwise = fromRelationInt $ (v `B.index` fromIntegral i) .&. relationMask
+
+-- | Get the indexes of members that satisfy the given relation predicate.
+getRelationsIndexes :: (MemberRelation -> Bool) -> ByteString -> [Int64]
+getRelationsIndexes p v = [i | i <- [0 .. fromIntegral (B.length v) - 1], p (getRelation i v)]
+
+-- | Set the relation status of a member at a given index in the relations vector.
+-- Expands the vector lazily if needed (padding with zeros for 'MRNew' relation).
+setRelation :: Int64 -> MemberRelation -> ByteString -> ByteString
+setRelation i r v
+  | i >= 0 = setRelations [(i, r)] v
+  | otherwise = v
+
+-- | Set multiple relations at once.
+-- Expands the vector lazily if needed (padding with zeros for 'MRNew' relation).
+setRelations :: [(Int64, MemberRelation)] -> ByteString -> ByteString
+setRelations [] v = v
+setRelations relations v =
+  let (fp, off, len) = toForeignPtr v
+      newLen = max len $ fromIntegral $ maximum (map fst relations) + 1
+   in unsafeCreate newLen $ \ptr -> do
+        withForeignPtr fp $ \vPtr -> copyBytes ptr (vPtr `plusPtr` off) len
+        when (newLen > len) $ fillBytes (ptr `plusPtr` len) 0 (newLen - len)
+        forM_ relations $ \(ix, r) -> when (ix >= 0) $ do
+          let i = fromIntegral ix
+          b <- peekByteOff ptr i
+          let b' = (b .&. complement relationMask) .|. toRelationInt r
+          pokeByteOff ptr i b'
+
+relationMask :: Word8
+relationMask = 0x07 -- reserving 3 bits

src/Simplex/Messaging/Agent/Store/SQLite.hs

@@ -42,18 +42,26 @@ module Simplex.Messaging.Agent.Store.SQLite
   )
 where
 
+import Control.Concurrent.MVar
+import Control.Concurrent.STM
+import Control.Exception (bracketOnError, onException, throwIO)
 import Control.Monad
 import Data.ByteArray (ScrubbedBytes)
 import qualified Data.ByteArray as BA
+import qualified Data.ByteString as B
 import Data.Functor (($>))
 import Data.IORef
 import Data.Maybe (fromMaybe)
+import Data.Int (Int64)
 import Data.Text (Text)
 import qualified Data.Text as T
 import Database.SQLite.Simple (Query (..))
 import qualified Database.SQLite.Simple as SQL
 import Database.SQLite.Simple.QQ (sql)
+import Database.SQLite3 (FuncArgs, FuncContext, funcArgInt64, funcResultBlob)
 import qualified Database.SQLite3 as SQLite3
+import Database.SQLite3.Direct (createAggregate)
+import Simplex.MemberRelations.MemberRelations (MemberRelation (..), fromRelationInt, setRelations)
 import Simplex.Messaging.Agent.Store.Migrations (DBMigrate (..), sharedMigrateSchema)
 import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.SQLite.Common
@@ -62,9 +70,6 @@ import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfig (..
 import Simplex.Messaging.Util (ifM, safeDecodeUtf8)
 import System.Directory (copyFile, createDirectoryIfMissing, doesFileExist)
 import System.FilePath (takeDirectory, takeFileName, (</>))
-import UnliftIO.Exception (bracketOnError, onException)
-import UnliftIO.MVar
-import UnliftIO.STM
 
 -- * SQLite Store implementation
 
@@ -109,16 +114,31 @@ connectDB path key track = do
   pure db
   where
     prepare db = do
-      let exec = SQLite3.exec $ SQL.connectionHandle $ DB.conn db
-      unless (BA.null key) . exec $ "PRAGMA key = " <> keyString key <> ";"
-      exec . fromQuery $
+      let db' = SQL.connectionHandle $ DB.conn db
+      unless (BA.null key) . SQLite3.exec db' $ "PRAGMA key = " <> keyString key <> ";"
+      SQLite3.exec db' . fromQuery $
         [sql|
           PRAGMA busy_timeout = 100;
           PRAGMA foreign_keys = ON;
           -- PRAGMA trusted_schema = OFF;
           PRAGMA secure_delete = ON;
           PRAGMA auto_vacuum = FULL;
         |]
+      createAggregate db' "build_relations_vector" (Just 2) [] buildRelationsVectorStep buildRelationsVectorFinal
+        >>= either (throwIO . userError . show) pure
+
+-- | Step function for build_relations_vector aggregate.
+-- Accumulates (idx, relation) pairs.
+buildRelationsVectorStep :: FuncContext -> FuncArgs -> [(Int64, MemberRelation)] -> IO [(Int64, MemberRelation)]
+buildRelationsVectorStep _ args acc = do
+  idx <- funcArgInt64 args 0
+  relation <- fromRelationInt . fromIntegral <$> funcArgInt64 args 1
+  pure $ (idx, relation) : acc
+
+-- | Final function for build_relations_vector aggregate.
+-- Builds the vector from accumulated pairs using setRelations.
+buildRelationsVectorFinal :: FuncContext -> [(Int64, MemberRelation)] -> IO ()
+buildRelationsVectorFinal ctx acc = funcResultBlob ctx $ setRelations acc B.empty
 
 closeDBStore :: DBStore -> IO ()
 closeDBStore st@DBStore {dbClosed} =