{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UnliftedFFITypes #-} {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# OPTIONS_HADDOCK not-home #-} {- HLINT ignore writeTable uoffsetFrom "Eta reduce" -} module FlatBuffers.Internal.Write where import Control.Monad.State.Strict import Data.Bits ( (.&.), complement ) import qualified Data.ByteString as BS import Data.ByteString.Builder ( Builder ) import qualified Data.ByteString.Builder as B import qualified Data.ByteString.Lazy as BSL import Data.Coerce ( coerce ) import Data.Int import qualified Data.List as L import qualified Data.Map.Strict as M import Data.MonoTraversable ( Element, MonoFoldable ) import qualified Data.MonoTraversable as Mono import Data.Monoid ( Sum(..) ) import Data.Semigroup ( Max(..) ) import Data.Text ( Text ) import qualified Data.Text.Array as A import qualified Data.Text.Encoding as T import qualified Data.Text.Internal as TI import Data.Word import FlatBuffers.Internal.Build import FlatBuffers.Internal.Constants import FlatBuffers.Internal.FileIdentifier ( FileIdentifier(unFileIdentifier), HasFileIdentifier(getFileIdentifier) ) import FlatBuffers.Internal.Types import Foreign.C.Types ( CSize(CSize) ) import GHC.Base ( ByteArray# ) import System.IO.Unsafe ( unsafeDupablePerformIO ) type BufferSize = Sum Int32 -- | The position of something in a buffer, expressed as the number of bytes counting from the end. type Position = Int32 data FBState = FBState { builder :: !Builder , bufferSize :: {-# UNPACK #-} !BufferSize , maxAlign :: {-# UNPACK #-} !(Max Alignment) , cache :: !(M.Map BSL.ByteString Position) } newtype WriteTableField = WriteTableField { unWriteTableField :: State FBState (FBState -> FBState) } -- | A struct to be written to a flatbuffer. newtype WriteStruct a = WriteStruct { buildStruct :: Builder } -- | A table to be written to a flatbuffer. newtype WriteTable a = WriteTable (State FBState Position) -- | A union to be written to a flatbuffer. data WriteUnion a = Some {-# UNPACK #-} !Word8 !(State FBState Position) | None -- | Serializes a flatbuffer table as a lazy `BSL.ByteString`. {-# INLINE encode #-} encode :: WriteTable a -> BSL.ByteString encode = encodeState (FBState mempty (Sum 0) (Max 1) mempty) {-# INLINE encodeState #-} encodeState :: FBState -> WriteTable a -> BSL.ByteString encodeState state (WriteTable writeTable) = B.toLazyByteString $ builder $ execState (do pos <- writeTable maxAlignment <- gets (getMax . maxAlign) modify' $ alignTo maxAlignment uoffsetSize modify' $ uoffsetFrom pos ) state -- | Serializes a flatbuffer table as a lazy `BSL.ByteString` and adds a File Identifier. {-# INLINE encodeWithFileIdentifier #-} encodeWithFileIdentifier :: forall a. HasFileIdentifier a => WriteTable a -> BSL.ByteString encodeWithFileIdentifier = encodeStateWithFileIdentifier (FBState mempty (Sum 0) (Max 1) mempty) (getFileIdentifier @a) {-# INLINE encodeStateWithFileIdentifier #-} encodeStateWithFileIdentifier :: FBState -> FileIdentifier -> WriteTable a -> BSL.ByteString encodeStateWithFileIdentifier state fi (WriteTable writeTable) = B.toLazyByteString $ builder $ execState (do pos <- writeTable maxAlignment <- gets (getMax . maxAlign) modify' $ alignTo maxAlignment (uoffsetSize + fileIdentifierSize) modify' $ writeFileIdentifier fi modify' $ uoffsetFrom pos ) state -- | Writes something (unaligned) to the buffer. {-# INLINE write #-} write :: Int32 -> Builder -> FBState -> FBState write bsize b fbs = fbs { builder = b <> builder fbs , bufferSize = bufferSize fbs <> Sum bsize } -- | Writes a 32-bit int (unaligned) to the buffer. {-# INLINE writeInt32 #-} writeInt32 :: Int32 -> FBState -> FBState writeInt32 n = write int32Size (B.int32LE n) {-# INLINE writeFileIdentifier #-} writeFileIdentifier :: FileIdentifier -> FBState -> FBState writeFileIdentifier fi = write fileIdentifierSize (B.byteString (unFileIdentifier fi)) {-# INLINE missing #-} missing :: WriteTableField missing = WriteTableField . pure $! id {-# INLINE deprecated #-} deprecated :: WriteTableField deprecated = missing {-# INLINE optional #-} optional :: (a -> WriteTableField) -> (Maybe a -> WriteTableField) optional = maybe missing {-# INLINE optionalDef #-} optionalDef :: Eq a => a -> (a -> WriteTableField) -> (Maybe a -> WriteTableField) optionalDef dflt write ma = case ma of Just a | a /= dflt -> write a _ -> missing {-# INLINE writeWord8TableField #-} writeWord8TableField :: Word8 -> WriteTableField writeWord8TableField n = WriteTableField . pure $! write word8Size (B.word8 n) . alignTo word8Size 0 {-# INLINE writeWord16TableField #-} writeWord16TableField :: Word16 -> WriteTableField writeWord16TableField n = WriteTableField . pure $! write word16Size (B.word16LE n) . alignTo word16Size 0 {-# INLINE writeWord32TableField #-} writeWord32TableField :: Word32 -> WriteTableField writeWord32TableField n = WriteTableField . pure $! write word32Size (B.word32LE n) . alignTo word32Size 0 {-# INLINE writeWord64TableField #-} writeWord64TableField :: Word64 -> WriteTableField writeWord64TableField n = WriteTableField . pure $! write word64Size (B.word64LE n) . alignTo word64Size 0 {-# INLINE writeInt8TableField #-} writeInt8TableField :: Int8 -> WriteTableField writeInt8TableField n = WriteTableField . pure $! write int8Size (B.int8 n) . alignTo int8Size 0 {-# INLINE writeInt16TableField #-} writeInt16TableField :: Int16 -> WriteTableField writeInt16TableField n = WriteTableField . pure $! write int16Size (B.int16LE n) . alignTo int16Size 0 {-# INLINE writeInt32TableField #-} writeInt32TableField :: Int32 -> WriteTableField writeInt32TableField n = WriteTableField . pure $! write int32Size (B.int32LE n) . alignTo int32Size 0 {-# INLINE writeInt64TableField #-} writeInt64TableField :: Int64 -> WriteTableField writeInt64TableField n = WriteTableField . pure $! write int64Size (B.int64LE n) . alignTo int64Size 0 {-# INLINE writeFloatTableField #-} writeFloatTableField :: Float -> WriteTableField writeFloatTableField n = WriteTableField . pure $! write floatSize (B.floatLE n) . alignTo floatSize 0 {-# INLINE writeDoubleTableField #-} writeDoubleTableField :: Double -> WriteTableField writeDoubleTableField n = WriteTableField . pure $! write doubleSize (B.doubleLE n) . alignTo doubleSize 0 {-# INLINE writeBoolTableField #-} writeBoolTableField :: Bool -> WriteTableField writeBoolTableField = writeWord8TableField . boolToWord8 {-# INLINE writeTextTableField #-} writeTextTableField :: Text -> WriteTableField writeTextTableField text = WriteTableField $ do modify' (writeInt32 len . encodeText . alignTo int32Size (len + 1)) uoffsetFromHere where len = utf8length text encodeText fbs = fbs -- strings must have a trailing zero { builder = T.encodeUtf8Builder text <> B.word8 0 <> builder fbs , bufferSize = Sum len <> Sum 1 <> bufferSize fbs } {-# INLINE writeTableTableField #-} writeTableTableField :: WriteTable a -> WriteTableField writeTableTableField (WriteTable writeTable) = WriteTableField $ do loc <- writeTable pure $! uoffsetFrom loc {-# INLINE writeStructTableField #-} writeStructTableField :: forall a. IsStruct a => WriteStruct a -> WriteTableField writeStructTableField (WriteStruct b) = writeStructTableField' (structAlignmentOf @a) (structSizeOf @a) b {-# INLINE writeStructTableField' #-} writeStructTableField' :: Alignment -> InlineSize -> Builder -> WriteTableField writeStructTableField' structAlignment structSize structBuilder = WriteTableField . pure $! writeStruct . alignTo structAlignment 0 where writeStruct fbs = fbs { builder = structBuilder <> builder fbs , bufferSize = bufferSize fbs <> Sum (fromIntegral @InlineSize @Int32 structSize) } {-# INLINE writeUnionTypesVectorTableField #-} writeUnionTypesVectorTableField :: WriteVector (WriteUnion a) -> WriteTableField writeUnionTypesVectorTableField (WriteVectorUnion tf _) = tf {-# INLINE writeUnionValuesVectorTableField #-} writeUnionValuesVectorTableField :: WriteVector (WriteUnion a) -> WriteTableField writeUnionValuesVectorTableField (WriteVectorUnion _ tf) = tf {-# INLINE writeUnionTypeTableField #-} writeUnionTypeTableField :: WriteUnion a -> WriteTableField writeUnionTypeTableField !wu = case wu of None -> missing Some unionType _ -> writeWord8TableField unionType {-# INLINE writeUnionValueTableField #-} writeUnionValueTableField :: WriteUnion a -> WriteTableField writeUnionValueTableField !wu = case wu of None -> missing Some _ unionValue -> writeTableTableField (WriteTable unionValue) -- | Constructs a missing union table field / vector element. {-# INLINE none #-} none :: WriteUnion a none = None {-# INLINE writeUnion #-} writeUnion :: Word8 -> WriteTable a -> WriteUnion b writeUnion n (WriteTable st) = Some n st {-# INLINE vtable #-} vtable :: [Word16] -> Word16 -> BSL.ByteString vtable fieldVOffsets tableSize = bytestring where vtableSize = voffsetSize + voffsetSize + voffsetSize * fromIntegral @Int @Word16 (L.length fieldVOffsets) bytestring = B.toLazyByteString ( B.word16LE vtableSize <> B.word16LE (coerce tableSize) <> foldMap (B.word16LE . coerce) fieldVOffsets ) {-# INLINE writeTable #-} writeTable :: [WriteTableField] -> WriteTable a writeTable fields = WriteTable $ do inlineFields <- sequence (coerce fields) -- table tableEnd <- gets (getSum . bufferSize) inlineFieldPositions <- forM inlineFields $ \f -> do before <- gets bufferSize modify' f after <- gets bufferSize if after == before then pure 0 else pure (getSum after) modify' $ alignTo soffsetSize 0 tableFieldsPosition <- gets (getSum . bufferSize) let tablePosition = tableFieldsPosition + soffsetSize -- Note: This might overflow if the table has too many fields let tableSize = fromIntegral @Int32 @Word16 $ tablePosition - tableEnd let fieldVOffsets = flip fmap inlineFieldPositions $ \case 0 -> 0 -- Note: This might overflow if the table has too many fields fieldPosition -> fromIntegral @Int32 @Word16 (tablePosition - fieldPosition) -- TODO: trim trailing 0 voffsets let newVtable = vtable fieldVOffsets tableSize let newVtableSize = fromIntegral @Int64 @Int32 (BSL.length newVtable) let newVtablePosition = tablePosition + newVtableSize map <- gets cache case M.insertLookupWithKey (\_k _new old -> old) newVtable newVtablePosition map of (Nothing, map') -> -- vtable, pointer to vtable, update the cache modify' (writeVtable map' newVtable newVtableSize . writeVtableSoffset newVtableSize) (Just oldVtablePosition, _) -> -- pointer to vtable modify' . writeInt32 . negate $ tablePosition - oldVtablePosition pure $! tablePosition where writeVtable newCache newVtable newVtableSize fbs = fbs { cache = newCache , builder = B.lazyByteString newVtable <> builder fbs , bufferSize = bufferSize fbs <> Sum newVtableSize } -- The vtable is located right before the table, so the offset -- between the table and the vtable is equal to the vtable size writeVtableSoffset newVtableSize = writeInt32 newVtableSize class WriteVectorElement a where -- | A vector to be written to a flatbuffer. data WriteVector a -- | Constructs a flatbuffers vector. -- -- If @n@ is larger than the length of @xs@, this will result in a malformed buffer. -- If @n@ is smaller than the length of @xs@, all elements of @xs@ will still be written to the buffer, -- but the client will only be able to read the first @n@ elements. -- -- Note: `fromMonoFoldable` asks for the collection's length to be passed in as an argument rather than use `Mono.olength` because: -- -- 1. `Mono.olength` is often O(n), and in some use cases there may be a better way to know the collection's length ahead of time. -- 2. Calling `Mono.olength` inside `fromMonoFoldable` can inhibit some fusions which would otherwise be possible. -- -- @since 0.2.0.0 -- Implementer's note: -- To elaborate on point 2., here's an example. -- This version of `fromMonoFoldable` that calls `Mono.olength` internally: -- -- > encodeUserIds' :: [User] -> BSL.ByteString -- > encodeUserIds' = encode . userIdsTable $ fromMonoFoldable (userId <$> users)) -- > -- > {-# INLINE fromMonoFoldable #-} -- > fromMonoFoldable xs = -- > let length = Mono.olength xs -- > buffer = foldr ... ... xs -- > in ... -- -- ...prevents `<$>` and `foldr` from being fused, and so it's much slower than when the length is passed in: -- -- > encodeUserIds :: [User] -> BSL.ByteString -- > encodeUserIds = encode . userIdsTable $ fromMonoFoldable (userId <$> users) (fromIntegral (Mono.olength users)) -- > -- > {-# INLINE fromMonoFoldable #-} -- > fromMonoFoldable xs length = -- > let buffer = foldr ... ... xs -- > in ... fromMonoFoldable :: (MonoFoldable mono, Element mono ~ a) => Int32 -- ^ @n@: the number of elements in @xs@ -> mono -- ^ @xs@: a collection -> WriteVector a -- | Convenience function, equivalent to: -- -- > fromMonoFoldable' xs = fromMonoFoldable (fromIntegral (olength xs)) xs -- -- In some cases it may be slower than using `fromMonoFoldable` directly. -- -- @since 0.2.0.0 {-# INLINE fromMonoFoldable' #-} fromMonoFoldable' :: (WriteVectorElement a, MonoFoldable mono, Element mono ~ a) => mono -> WriteVector a fromMonoFoldable' xs = fromMonoFoldable (fromIntegral $ Mono.olength xs) xs -- | `fromMonoFoldable` specialized to list fromList :: WriteVectorElement a => Int32 -> [a] -> WriteVector a fromList = fromMonoFoldable -- | `fromMonoFoldable'` specialized to list fromList' :: WriteVectorElement a => [a] -> WriteVector a fromList' = fromMonoFoldable' -- | Creates a flatbuffers vector with a single element singleton :: WriteVectorElement a => a -> WriteVector a singleton a = fromList 1 [a] -- | Creates an empty flatbuffers vector empty :: WriteVectorElement a => WriteVector a empty = fromList 0 [] newtype FromFoldable f a = FromFoldable (f a) deriving newtype Foldable type instance Element (FromFoldable f a) = a instance Foldable f => MonoFoldable (FromFoldable f a) -- | `fromMonoFoldable` for types that implement `Foldable` but not `MonoFoldable`. fromFoldable :: (WriteVectorElement a, Foldable f) => Int32 -> f a -> WriteVector a fromFoldable n = fromMonoFoldable n . FromFoldable -- | `fromMonoFoldable'` for types that implement `Foldable` but not `MonoFoldable`. fromFoldable' :: (WriteVectorElement a, Foldable f) => f a -> WriteVector a fromFoldable' = fromMonoFoldable' . FromFoldable -- | Efficiently creates a vector from a `BS.ByteString`. -- Large `BS.ByteString`s are inserted directly, but small ones are copied to ensure that the generated chunks are large on average. -- -- @since 0.2.0.0 fromByteString :: BS.ByteString -> WriteVector Word8 fromByteString bs = WriteVectorWord8 . WriteTableField $ do modify' $! writeInt32 len . writeByteString . alignTo int32Size len uoffsetFromHere where len = fromIntegral @Int @Int32 (BS.length bs) writeByteString fbs = fbs { builder = B.byteString bs <> builder fbs , bufferSize = bufferSize fbs <> Sum len } -- | Efficiently creates a vector from a lazy `BSL.ByteString`. -- Large chunks of the `BSL.ByteString` are inserted directly, but small ones are copied to ensure that the generated chunks are large on average. -- -- @since 0.2.0.0 fromLazyByteString :: BSL.ByteString -> WriteVector Word8 fromLazyByteString bs = WriteVectorWord8 . WriteTableField $ do modify' $! writeInt32 len . writeByteString . alignTo int32Size len uoffsetFromHere where len = fromIntegral @Int64 @Int32 (BSL.length bs) writeByteString fbs = fbs { builder = B.lazyByteString bs <> builder fbs , bufferSize = bufferSize fbs <> Sum len } {-# INLINE inlineVector #-} inlineVector :: (MonoFoldable mono, Element mono ~ a) => (a -> Builder) -> Alignment -> InlineSize -> Int32 -> mono -> WriteTableField inlineVector build elemAlignment elemSize elemCount elems = WriteTableField $ do modify' $! writeInt32 elemCount . writeVec . alignTo (coerce elemAlignment `max` int32Size) vecByteLength uoffsetFromHere where vecByteLength = elemCount * fromIntegral @InlineSize @Int32 elemSize vecBuilder = Mono.ofoldr (\a b -> build a <> b) mempty elems writeVec fbs = fbs { builder = vecBuilder <> builder fbs , bufferSize = bufferSize fbs <> Sum vecByteLength } instance WriteVectorElement Word8 where newtype WriteVector Word8 = WriteVectorWord8 { writeVectorWord8TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Word8) => Int32 -> mono -> WriteVector Word8 fromMonoFoldable n = WriteVectorWord8 . inlineVector B.word8 word8Size word8Size n instance WriteVectorElement Word16 where newtype WriteVector Word16 = WriteVectorWord16 { writeVectorWord16TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Word16) => Int32 -> mono -> WriteVector Word16 fromMonoFoldable n = WriteVectorWord16 . inlineVector B.word16LE word16Size word16Size n instance WriteVectorElement Word32 where newtype WriteVector Word32 = WriteVectorWord32 { writeVectorWord32TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Word32) => Int32 -> mono -> WriteVector Word32 fromMonoFoldable n = WriteVectorWord32 . inlineVector B.word32LE word32Size word32Size n instance WriteVectorElement Word64 where newtype WriteVector Word64 = WriteVectorWord64 { writeVectorWord64TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Word64) => Int32 -> mono -> WriteVector Word64 fromMonoFoldable n = WriteVectorWord64 . inlineVector B.word64LE word64Size word64Size n instance WriteVectorElement Int8 where newtype WriteVector Int8 = WriteVectorInt8 { writeVectorInt8TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Int8) => Int32 -> mono -> WriteVector Int8 fromMonoFoldable n = WriteVectorInt8 . inlineVector B.int8 int8Size int8Size n instance WriteVectorElement Int16 where newtype WriteVector Int16 = WriteVectorInt16 { writeVectorInt16TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Int16) => Int32 -> mono -> WriteVector Int16 fromMonoFoldable n = WriteVectorInt16 . inlineVector B.int16LE int16Size int16Size n instance WriteVectorElement Int32 where newtype WriteVector Int32 = WriteVectorInt32 { writeVectorInt32TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Int32) => Int32 -> mono -> WriteVector Int32 fromMonoFoldable n = WriteVectorInt32 . inlineVector B.int32LE int32Size int32Size n instance WriteVectorElement Int64 where newtype WriteVector Int64 = WriteVectorInt64 { writeVectorInt64TableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Int64) => Int32 -> mono -> WriteVector Int64 fromMonoFoldable n = WriteVectorInt64 . inlineVector B.int64LE int64Size int64Size n instance WriteVectorElement Float where newtype WriteVector Float = WriteVectorFloat { writeVectorFloatTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Float) => Int32 -> mono -> WriteVector Float fromMonoFoldable n = WriteVectorFloat . inlineVector B.floatLE floatSize floatSize n instance WriteVectorElement Double where newtype WriteVector Double = WriteVectorDouble { writeVectorDoubleTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Double) => Int32 -> mono -> WriteVector Double fromMonoFoldable n = WriteVectorDouble . inlineVector B.doubleLE doubleSize doubleSize n instance WriteVectorElement Bool where newtype WriteVector Bool = WriteVectorBool { writeVectorBoolTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Bool) => Int32 -> mono -> WriteVector Bool fromMonoFoldable n = WriteVectorBool . inlineVector (B.word8 . boolToWord8) word8Size word8Size n instance IsStruct a => WriteVectorElement (WriteStruct a) where newtype WriteVector (WriteStruct a) = WriteVectorStruct { writeVectorStructTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ WriteStruct a) => Int32 -> mono -> WriteVector (WriteStruct a) fromMonoFoldable n = WriteVectorStruct . inlineVector coerce (structAlignmentOf @a) (structSizeOf @a) n data TextInfos = TextInfos ![TextInfo] {-# UNPACK #-} !BufferSize data TextInfo = TextInfo { tiText :: !Text , tiUtf8len :: {-# UNPACK #-} !Int32 , tiPadding :: {-# UNPACK #-} !Int32 , tiPosition :: {-# UNPACK #-} !Position } data OffsetInfo = OffsetInfo { oiIndex :: {-# UNPACK #-} !Int32 , oiOffsets :: ![Int32] } instance WriteVectorElement Text where newtype WriteVector Text = WriteVectorText { writeVectorTextTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ Text) => Int32 -> mono -> WriteVector Text fromMonoFoldable elemCount texts = WriteVectorText . WriteTableField $ do modify' $ \fbs -> let (builder2, bsize2) = writeVectorSizePrefix . writeOffsets . align . writeStrings $ (builder fbs, bufferSize fbs) in fbs { builder = builder2 , bufferSize = bsize2 , maxAlign = maxAlign fbs <> Max int32Size } uoffsetFromHere where writeStrings :: (Builder, BufferSize) -> (Builder, BufferSize, [TextInfo]) writeStrings (builder1, bsize1) = -- Collect info about the strings. -- NOTE: this loop *could* be merged with the one below, but -- we have loops dedicated to merging Builders to avoid wrapping Builders in data structures. -- See "Performance tips": http://hackage.haskell.org/package/fast-builder-0.1.0.1/docs/Data-ByteString-FastBuilder.html let TextInfos textInfos bsize2 = Mono.ofoldr (\t (TextInfos infos bsize) -> let textLength = utf8length t padding = calcPadding 4 (textLength + 1) bsize newBsize = bsize <> Sum (padding + textLength + 1 + 4) in TextInfos (TextInfo t textLength padding (getSum newBsize) : infos) newBsize ) (TextInfos [] bsize1) texts builder2 = foldr (\(TextInfo t tlength padding _) b -> B.int32LE tlength <> T.encodeUtf8Builder t <> B.word8 0 -- strings must have a trailing zero <> buildPadding padding <> b ) mempty textInfos in (builder2 <> builder1, bsize2, textInfos) align :: (Builder, BufferSize, [TextInfo]) -> (Builder, BufferSize, [TextInfo]) align (builder1, bsize1, textInfos) = let vectorPadding = calcPadding int32Size 0 bsize1 bsize2 = bsize1 <> Sum vectorPadding builder2 = buildPadding vectorPadding in (builder2 <> builder1, bsize2, textInfos) writeOffsets :: (Builder, BufferSize, [TextInfo]) -> (Builder, BufferSize) writeOffsets (builder1, bsize1, textInfos) = let OffsetInfo _ offsets = foldr (\(TextInfo _ _ _ position) (OffsetInfo ix os) -> OffsetInfo (ix + 1) (getSum bsize1 + (ix * 4) + 4 - position : os) ) (OffsetInfo 0 []) textInfos bsize2 = bsize1 <> Sum (elemCount * 4) builder2 = foldr (\o b -> B.int32LE o <> b) mempty offsets in (builder2 <> builder1, bsize2) writeVectorSizePrefix :: (Builder, BufferSize) -> (Builder, BufferSize) writeVectorSizePrefix (builder1, bsize1) = (B.int32LE elemCount <> builder1, bsize1 + int32Size) data TableInfo = TableInfo { tiState :: !FBState , tiTablePositions :: ![Position] } instance WriteVectorElement (WriteTable a) where newtype WriteVector (WriteTable a) = WriteVectorTable { writeVectorTableTableField :: WriteTableField } {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ WriteTable a) => Int32 -> mono -> WriteVector (WriteTable a) fromMonoFoldable elemCount tables = WriteVectorTable . WriteTableField $ do fbs1 <- get let !(TableInfo fbs2 positions) = Mono.ofoldr (\(WriteTable writeTable) (TableInfo fbs positions) -> let (pos, fbs') = runState writeTable fbs in TableInfo fbs' (pos : positions) ) (TableInfo fbs1 []) tables put $! alignTo int32Size 0 fbs2 -- Write offsets bsize <- gets (getSum . bufferSize) let OffsetInfo _ offsets = foldr (\position (OffsetInfo ix os) -> OffsetInfo (ix + 1) (bsize + (ix * 4) + 4 - position : os) ) (OffsetInfo 0 []) positions coerce $ fromMonoFoldable elemCount offsets data Vecs a = Vecs ![Word8] ![Maybe (State FBState Position)] data UnionTableInfo = UnionTableInfo { utiState :: !FBState , utiTablePositions :: ![Maybe Position] } instance WriteVectorElement (WriteUnion a) where data WriteVector (WriteUnion a) = WriteVectorUnion !WriteTableField !WriteTableField {-# INLINE fromMonoFoldable #-} fromMonoFoldable :: (MonoFoldable mono, Element mono ~ WriteUnion a) => Int32 -> mono -> WriteVector (WriteUnion a) fromMonoFoldable elemCount unions = let Vecs types values = Mono.ofoldr go (Vecs [] []) unions go writeUnion (Vecs types values) = case writeUnion of None -> Vecs (0 : types) (Nothing : values) Some typ val -> Vecs (typ : types) (Just val : values) writeUnionTables :: WriteTableField writeUnionTables = WriteTableField $ do fbs1 <- get let !(UnionTableInfo fbs2 positions) = foldr (\unionTableOpt (UnionTableInfo fbs positions) -> case unionTableOpt of Just t -> let (pos, fbs') = runState t fbs in UnionTableInfo fbs' (Just pos : positions) Nothing -> UnionTableInfo fbs (Nothing : positions) ) (UnionTableInfo fbs1 []) values put $! alignTo int32Size 0 fbs2 -- Write offsets bsize <- gets (getSum . bufferSize) let OffsetInfo _ offsets = foldr (\positionOpt (OffsetInfo ix os) -> let offset = case positionOpt of Just position -> bsize + (ix * 4) + 4 - position Nothing -> 0 in OffsetInfo (ix + 1) (offset : os) ) (OffsetInfo 0 []) positions coerce $ fromMonoFoldable elemCount offsets in WriteVectorUnion (coerce $ fromMonoFoldable elemCount types) writeUnionTables -- | Calculate how much 0-padding is needed so that, after writing @additionalBytes@, -- the buffer becomes aligned to @n@ bytes. {-# INLINE calcPadding #-} calcPadding :: Alignment {- ^ n -} -> Int32 {- ^ additionalBytes -} -> BufferSize -> Int32 calcPadding !n !additionalBytes (Sum size) = (complement (size + additionalBytes) + 1) .&. (fromIntegral n - 1) -- | Add enough 0-padding so that the buffer becomes aligned to @n@ after writing @additionalBytes@. {-# INLINE alignTo #-} alignTo :: Alignment{- ^ n -} -> Int32 {- ^ additionalBytes -} -> FBState -> FBState alignTo !n !additionalBytes fbs@(FBState b bsize ma cache) = if padding == 0 then fbs { maxAlign = ma <> coerce n } else FBState (buildPadding padding <> b) (bsize <> Sum padding) (ma <> coerce n) cache where padding = calcPadding n additionalBytes bsize {-# INLINE uoffsetFromHere #-} uoffsetFromHere :: State FBState (FBState -> FBState) uoffsetFromHere = gets (uoffsetFrom . coerce . bufferSize) {-# INLINE uoffsetFrom #-} uoffsetFrom :: Position -> FBState -> FBState uoffsetFrom pos = writeUOffset . align where align fbs = alignTo int32Size 0 fbs writeUOffset fbs = let currentPos = coerce bufferSize fbs in writeInt32 (currentPos - pos + uoffsetSize) fbs {-# INLINE utf8length #-} utf8length :: Text -> Int32 utf8length (TI.Text arr off len) | len == 0 = 0 | otherwise = unsafeDupablePerformIO $ c_length_utf8 (A.aBA arr) (fromIntegral off) (fromIntegral len) foreign import ccall unsafe "_hs_text_length_utf8" c_length_utf8 :: ByteArray# -> CSize -> CSize -> IO Int32