import BN from 'bn.js' ; import { blob , offset , seq , struct , u32 , u8 , union } from 'buffer-layout' ; import { publicKeyLayout , setLayoutDecoder , u128 , u64 , zeros } from './layout' ; import { PublicKey } from '@solana/web3.js' ; const SLAB_HEADER_LAYOUT = struct ( [ u32 ( 'bumpIndex' ) , zeros ( 4 ) , u32 ( 'freeListLen' ) , zeros ( 4 ) , u32 ( 'freeListHead' ) , u32 ( 'root' ) , u32 ( 'leafCount' ) , zeros ( 4 ) , ] , 'header' ) ; const SLAB_NODE_LAYOUT = union ( u32 ( 'tag' ) , blob ( 68 ) , 'node' ) ; SLAB_NODE_LAYOUT . addVariant ( 0 , struct ( [ ] ) , 'uninitialized' ) ; SLAB_NODE_LAYOUT . addVariant ( 1 , struct ( [ u32 ( 'prefixLen' ) , u128 ( 'key' ) , seq ( u32 ( ) , 2 , 'children' ) , ] ) , 'innerNode' ) ; SLAB_NODE_LAYOUT . addVariant ( 2 , struct ( [ u8 ( 'ownerSlot' ) , u8 ( 'feeTier' ) , blob ( 2 ) , u128 ( 'key' ) , publicKeyLayout ( 'owner' ) , u64 ( 'quantity' ) , u64 ( 'clientOrderId' ) , ] ) , 'leafNode' ) ; SLAB_NODE_LAYOUT . addVariant ( 3 , struct ( [ u32 ( 'next' ) ] ) , 'freeNode' ) ; SLAB_NODE_LAYOUT . addVariant ( 4 , struct ( [ ] ) , 'lastFreeNode' ) ; export const SLAB_LAYOUT = struct ( [ SLAB_HEADER_LAYOUT , seq ( SLAB_NODE_LAYOUT , offset ( SLAB_HEADER_LAYOUT . layoutFor ( 'bumpIndex' ) , SLAB_HEADER_LAYOUT . offsetOf ( 'bumpIndex' ) - SLAB_HEADER_LAYOUT . span ) , 'nodes' ) , ] ) ; export class Slab { private header : any ; private nodes : any ; constructor ( header , nodes ) { this . header = header ; this . nodes = nodes ; } static decode ( buffer ) { return SLAB_LAYOUT . decode ( buffer ) ; } get ( searchKey ) { if ( this . header . leafCount === 0 ) { return null ; } if ( ! ( searchKey instanceof BN ) ) { searchKey = new BN ( searchKey ) ; } let index = this . header . root ; while ( true ) { const { leafNode , innerNode } = this . nodes [ index ] ; if ( leafNode ) { if ( leafNode . key . eq ( searchKey ) ) { return leafNode ; } return null ; } else if ( innerNode ) { if ( ! innerNode . key . xor ( searchKey ) . iushrn ( 128 - innerNode . prefixLen ) . isZero ( ) ) { return null ; } index = innerNode . children [ searchKey . testn ( 128 - innerNode . prefixLen - 1 ) ? 1 : 0 ] ; } else { throw new Error ( 'Invalid slab' ) ; } } } [ Symbol . iterator ] ( ) { return this . items ( false ) ; } * items ( descending = false ) { if ( this . header . leafCount === 0 ) { return ; } const stack = [ this . header . root ] ; while ( stack . length > 0 ) { const index = stack . pop ( ) ; const { leafNode , innerNode } = this . nodes [ index ] ; if ( leafNode ) { yield leafNode ; } else if ( innerNode ) { if ( descending ) { stack . push ( innerNode . children [ 0 ] , innerNode . children [ 1 ] ) ; } else { stack . push ( innerNode . children [ 1 ] , innerNode . children [ 0 ] ) ; } } } } } setLayoutDecoder ( SLAB_LAYOUT , ( { header , nodes } ) => new Slab ( header , nodes ) ) ;