pyaver.slab

  1from dataclasses import dataclass
  2from typing import Iterable, List, NamedTuple, Optional
  3from construct import ListContainer
  4from solana.publickey import PublicKey
  5from .layouts import SLAB_LAYOUT, NodeType, CALLBACK_INFO_LEN
  6
  7class Callback(NamedTuple):
  8    user_market: PublicKey
  9    fee_tier: int
 10
 11class SlabHeader(NamedTuple):
 12    account_tag: int
 13    bump_index: int
 14    free_list_len: int
 15    free_list_head: int
 16    callback_memory_offset: int
 17    callback_free_list_len: int
 18    callback_free_list_head: int
 19    callback_bump_index: int
 20    root_node: int
 21    leaf_count: int
 22    market_address: PublicKey
 23
 24
 25# Used as dummy value for SlabNode#next.
 26NONE_NEXT = -1
 27
 28
 29# UninitializedNode, FreeNode and LastFreeNode all maps to this class.
 30@dataclass(frozen=True)
 31class SlabNode:
 32    is_initialized: bool
 33    next: int
 34
 35@dataclass(frozen=True)
 36class SlabLeafNode(SlabNode):
 37    key: int
 38    callback_info_pt: int
 39    base_quantity: int
 40    user_market: PublicKey
 41    fee_tier: int
 42
 43@dataclass(frozen=True)
 44class SlabInnerNode(SlabNode):
 45    prefix_len: int
 46    key: int
 47    children: List[int]
 48
 49
 50class Slab:
 51    """
 52    Slab object
 53
 54    A Slab is one side of the orderbook (bids or asks), which contains all of the open orders for side and outcome in a given market
 55    """
 56    def __init__(self, header: SlabHeader, nodes: List[SlabNode]): 
 57        self._header: SlabHeader = header
 58        self._nodes: List[SlabNode] = nodes
 59
 60    @staticmethod
 61    def __build(nodes: ListContainer, buffer: bytes) -> List[SlabNode]:
 62        res: List[SlabNode] = []
 63        for construct_node in nodes:
 64            node_type = construct_node.tag
 65            node = construct_node.node
 66            if node_type == NodeType.UNINTIALIZED:
 67                res.append(SlabNode(is_initialized=False, next=-1))
 68            elif node_type == NodeType.LEAF_NODE:
 69                res.append(
 70                    SlabLeafNode(
 71                        key = int.from_bytes(node.key, "little"),
 72                        callback_info_pt = node.callback_info_pt,
 73                        base_quantity = node.base_quantity,
 74                        is_initialized = True,
 75                        next = NONE_NEXT,
 76                        user_market = PublicKey(buffer[node.callback_info_pt:node.callback_info_pt+CALLBACK_INFO_LEN-1]),
 77                        fee_tier = int.from_bytes(buffer[node.callback_info_pt+CALLBACK_INFO_LEN-1:node.callback_info_pt+CALLBACK_INFO_LEN], "little"),
 78                    )
 79                )
 80            elif node_type == NodeType.INNER_NODE:
 81                res.append(
 82                    SlabInnerNode(
 83                        prefix_len = node.prefix_len,
 84                        key = int.from_bytes(node.key, "little"),
 85                        children = node.children,
 86                        is_initialized = True,
 87                        next = NONE_NEXT,
 88                    )
 89                )
 90            elif node_type == NodeType.FREE_NODE:
 91                res.append(
 92                    SlabNode(
 93                        is_initialized=True,
 94                        next=node.next
 95                    )
 96                )
 97            elif node_type == NodeType.LAST_FREE_NODE:
 98                res.append(
 99                    SlabNode(
100                        is_initialized=True,
101                        next=NONE_NEXT
102                    )
103                )
104            else:
105                raise RuntimeError("Unrecognized node type" + node.tag)
106        return res        
107
108    @staticmethod
109    def from_bytes(buffer: bytes):
110        """
111        Parses raw onchain data to Slab object
112
113        Args:
114            buffer (bytes): Raw bytes coming from onchain
115
116        Returns:
117            Slab: Slab object
118        """
119        parsed_slab = SLAB_LAYOUT.parse(buffer)
120        header = parsed_slab.header
121        nodes = parsed_slab.nodes
122        return Slab(
123            SlabHeader(
124                account_tag=header.account_tag,
125                bump_index=header.bump_index,
126                free_list_len=header.free_list_len,
127                free_list_head=header.free_list_head,
128                callback_memory_offset=header.callback_memory_offset,
129                callback_free_list_len=header.callback_free_list_len,
130                callback_free_list_head=header.callback_free_list_head,
131                callback_bump_index=header.callback_bump_index,
132                root_node=header.root_node,
133                leaf_count=header.leaf_count,
134                market_address=PublicKey(header.market_address),
135            ),
136            Slab.__build(nodes, buffer),
137        )
138
139    def get(self, search_key: int) -> Optional[SlabLeafNode]:
140        if self._header.leaf_count == 0:
141            return None
142        index: int = self._header.root_node
143        while True:
144            node: SlabNode = self._nodes[index]
145            if isinstance(node, SlabLeafNode):  # pylint: disable=no-else-return
146                return node if node.key == search_key else None
147            elif isinstance(node, SlabInnerNode):
148                if (node.key ^ search_key) >> (128 - node.prefix_len) != 0:
149                    return None
150                # Check if the n-th bit (start from the least significant, i.e. rightmost) of the key is set
151                index = node.children[(search_key >> (128 - node.prefix_len - 1)) & 1]
152            else:
153                raise RuntimeError("Should not go here! Node type not recognize.")
154
155    def __iter__(self) -> Iterable[SlabLeafNode]:
156        return self.items(False)
157
158    def items(self, descending=False) -> Iterable[SlabLeafNode]:
159        """
160        Depth first traversal of the Binary Tree of orders in the Slab
161
162        Args:
163            descending (bool, optional): Decides if the price should descending or not. Defaults to False.
164
165        Raises:
166            RuntimeError: Neither of leaf node or tree node!
167
168        Returns:
169            Iterable[SlabLeafNode]: SlabLeafNode object
170        """
171        if self._header.leaf_count == 0:
172            return
173        stack = [self._header.root_node]
174        while stack:
175            index = stack.pop()
176            node: SlabNode = self._nodes[index]
177            if isinstance(node, SlabLeafNode):
178                yield node
179            elif isinstance(node, SlabInnerNode):
180                if descending:
181                    stack.append(node.children[0])
182                    stack.append(node.children[1])
183                else:
184                    stack.append(node.children[1])
185                    stack.append(node.children[0])
186            else:
187                raise RuntimeError("Neither of leaf node or tree node!")