Inside Bitcoin

251                return _BinNode(node0.ins, node0.arg1 + (1 << (node0.arg1.bit_length() - 1)),
252                                node0.arg2)
253            return _BinNode(_Instruction.MATCH, 2, node0)
254        return _BinNode(_Instruction.JUMP, node0, node1)
255
256    @staticmethod
257    def make_default(val: int, sub: "_BinNode") -> "_BinNode":
258        """
259        Construct a _BinNode that corresponds to the specified subprogram, with the specified
260        default value. It exploits shortcuts that are possible in the encoding, and will use
261        either a DEFAULT or a RETURN instruction."""
262        assert val is not None and val > 0
263        if sub.ins == _Instruction.END:
264            return _BinNode(_Instruction.RETURN, val)
265        if sub.ins in (_Instruction.RETURN, _Instruction.DEFAULT):
266            return sub
267        return _BinNode(_Instruction.DEFAULT, val, sub)
268
269@total_ordering
270class ASMap:
271    """
272    A class whose objects represent a mapping from subnets to ASNs.
273
274    Internally the mapping is stored as a binary trie, but can be converted
275    from/to a list of ASNEntry objects, and from/to the binary asmap file format.
276
277    In the trie representation, nodes are represented as bare lists for efficiency
278    and ease of manipulation:
279    - [0] means an unassigned subnet (no ASN mapping for it is present)
280    - [int] means a subnet mapped entirely to the specified ASN.
281    - [node,node] means a subnet whose lower half and upper half have different
282    -             mappings, represented by new trie nodes.
283    """
284
285    def update(self, prefix: list[bool], asn: int) -> None:
286        """Update this ASMap object to map prefix to the specified asn."""
287        assert asn == 0 or _CODER_ASN.can_encode(asn)
288
289        def recurse(node: list, offset: int) -> None:
290            if offset == len(prefix):
291                # Reached the end of prefix; overwrite this node.
292                node.clear()
293                node.append(asn)
294                return
295            if len(node) == 1:
296                # Need to descend into a leaf node; split it up.
297                oldasn = node[0]
298                node.clear()
299                node.append([oldasn])
300                node.append([oldasn])
301            # Descend into the node.
302            recurse(node[prefix[offset]], offset + 1)
303            # If the result is two identical leaf children, merge them.
304            if len(node[0]) == 1 and len(node[1]) == 1 and node[0] == node[1]:
305                oldasn = node[0][0]
306                node.clear()
307                node.append(oldasn)
308        recurse(self._trie, 0)
309
310    def update_multi(self, entries: list[tuple[list[bool], int]]) -> None:
311        """Apply multiple update operations, where longer prefixes take precedence."""
312        entries.sort(key=lambda entry: len(entry[0]))
313        for prefix, asn in entries:
314            self.update(prefix, asn)
315
316    def _set_trie(self, trie) -> None:
317        """Set trie directly. Internal use only."""
318        def recurse(node: list) -> None:
319            if len(node) < 2:
320                return
321            recurse(node[0])
322            recurse(node[1])
323            if len(node[0]) == 2:
324                return
325            if node[0] == node[1]:
326                if len(node[0]) == 0:
327                    node.clear()
328                else:
329                    asn = node[0][0]
330                    node.clear()
331                    node.append(asn)
332        recurse(trie)
333        self._trie = trie
334
335    def __init__(self, entries: Optional[Iterable[ASNEntry]] = None) -> None:
336        """Construct an ASMap object from an optional list of entries."""
337        self._trie = [0]
338        if entries is not None:
339            def entry_key(entry):
340                """Sort function that places shorter prefixes first."""
341                prefix, asn = entry
342                return len(prefix), prefix, asn
343            for prefix, asn in sorted(entries, key=entry_key):
344                self.update(prefix, asn)
345
346    def lookup(self, prefix: list[bool]) -> Optional[int]:
347        """Look up a prefix. Returns ASN, or 0 if unassigned, or None if indeterminate."""
348        node = self._trie
349        for bit in prefix:
350            if len(node) == 1:
351                break
352            node = node[bit]
353        if len(node) == 1:
354            return node[0]
355        return None
356
357    def _to_entries_flat(self, fill: bool = False) -> list[ASNEntry]:
358        """Convert an ASMap object to a list of non-overlapping (prefix, asn) objects."""
359        prefix : list[bool] = []
360
361        def recurse(node: list) -> list[ASNEntry]:
362            ret = []
363            if len(node) == 1:
364                if node[0] > 0:
365                    ret = [(list(prefix), node[0])]
366            elif len(node) == 2:
367                prefix.append(False)
368                ret = recurse(node[0])
369                prefix[-1] = True
370                ret += recurse(node[1])
371                prefix.pop()
372                if fill and len(ret) > 1:
373                    asns = set(x[1] for x in ret)
374                    if len(asns) == 1:
375                        ret = [(list(prefix), list(asns)[0])]
376            return ret
377        return recurse(self._trie)
378
379    def _to_entries_minimal(self, fill: bool = False) -> list[ASNEntry]:
380        """Convert a trie to a minimal list of ASNEntry objects, exploiting overlap."""
381        prefix : list[bool] = []
382
383        def recurse(node: list) -> (tuple[dict[Optional[int], list[ASNEntry]], bool]):
384            if len(node) == 1 and node[0] == 0:
385                return {None if fill else 0: []}, True
386            if len(node) == 1:
387                return {node[0]: [], None: [(list(prefix), node[0])]}, False
388            ret: dict[Optional[int], list[ASNEntry]] = {}
389            prefix.append(False)
390            left, lhole = recurse(node[0])
391            prefix[-1] = True
392            right, rhole = recurse(node[1])
393            prefix.pop()
394            hole = not fill and (lhole or rhole)
395            def candidate(ctx: Optional[int], res0: Optional[list[ASNEntry]],
396                    res1: Optional[list[ASNEntry]]):
397                if res0 is not None and res1 is not None:
398                    if ctx not in ret or len(res0) + len(res1) < len(ret[ctx]):
399                        ret[ctx] = res0 + res1
400            for ctx in set(left) | set(right):
401                candidate(ctx, left.get(ctx), right.get(ctx))
402                candidate(ctx, left.get(None), right.get(ctx))
403                candidate(ctx, left.get(ctx), right.get(None))
404            if not hole:
405                for ctx in list(ret):
406                    if ctx is not None:
407                        candidate(None, [(list(prefix), ctx)], ret[ctx])
408            if None in ret:
409                ret = {ctx:entries for ctx, entries in ret.items()
410                       if ctx is None or len(entries) < len(ret[None])}
411            if hole:
412                ret = {ctx:entries for ctx, entries in ret.items() if ctx is None or ctx == 0}
413            return ret, hole
414        res, _ = recurse(self._trie)
415        return res[0] if 0 in res else res[None]
416
417    def __str__(self) -> str:
418        """Convert this ASMap object to a string containing Python code constructing it."""
419        return f"ASMap({self._trie})"
420
421    def to_entries(self, overlapping: bool = True, fill: bool = False) -> list[ASNEntry]:
422        """
423        Convert the mappings in this ASMap object to a list of ASNEntry objects.
424
425        Arguments:
426            overlapping: Permit the subnets in the resulting ASNEntry to overlap.
427                         Setting this can result in a shorter list.
428            fill:        Permit the resulting ASNEntry objects to cover subnets that
429                         are unassigned in this ASMap object. Setting this can
430                         result in a shorter list.
431        """
432        if overlapping:
433            return self._to_entries_minimal(fill)
434        return self._to_entries_flat(fill)
435
436    @staticmethod
437    def from_random(num_leaves: int = 10, max_asn: int = 6,
438                    unassigned_prob: float = 0.5) -> "ASMap":
439        """
440        Construct a random ASMap object, with specified:
441         - Number of leaves in its trie (at least 1)
442         - Maximum ASN value (at least 1)
443         - Probability for leaf nodes to be unassigned
444
445        The number of leaves in the resulting object may be less than what is
446        requested. This method is mostly intended for testing.
447        """
448        assert num_leaves >= 1
449        assert max_asn >= 1 or unassigned_prob == 1
450        assert _CODER_ASN.can_encode(max_asn)
451        assert 0.0 <= unassigned_prob <= 1.0
452        trie: list = []
453        leaves = [trie]
454        ret = ASMap()
455        for i in range(1, num_leaves):
456            idx = random.randrange(i)
457            leaf = leaves[idx]
458            lastleaf = leaves.pop()
459            if idx + 1 < i:
460                leaves[idx] = lastleaf
461            leaf.append([])
462            leaf.append([])
463            leaves.append(leaf[0])
464            leaves.append(leaf[1])
465        for leaf in leaves:
466            if random.random() >= unassigned_prob:
467                leaf.append(random.randrange(1, max_asn + 1))
468            else:
469                leaf.append(0)
470        #pylint: disable=protected-access
471        ret._set_trie(trie)
472        return ret
473
474    def _to_binnode(self, fill: bool = False) -> _BinNode:
475        """Convert a trie to a _BinNode object."""
476        def recurse(node: list) -> tuple[dict[Optional[int], _BinNode], bool]:
477            if len(node) == 1 and node[0] == 0:
478                return {(None if fill else 0): _BinNode.make_end()}, True
479            if len(node) == 1:
480                return {None: _BinNode.make_leaf(node[0]), node[0]: _BinNode.make_end()}, False
481            ret: dict[Optional[int], _BinNode] = {}
482            left, lhole = recurse(node[0])
483            right, rhole = recurse(node[1])
484            hole = (lhole or rhole) and not fill
485
486            def candidate(ctx: Optional[int], arg1, arg2, func: Callable):
487                if arg1 is not None and arg2 is not None:
488                    cand = func(arg1, arg2)
489                    if ctx not in ret or cand.size < ret[ctx].size:
490                        ret[ctx] = cand
491
492            union = set(left) | set(right)
493            sorted_union = sorted(union, key=lambda x: (x is None, x))
494            for ctx in sorted_union:
495                candidate(ctx, left.get(ctx), right.get(ctx), _BinNode.make_branch)
496                candidate(ctx, left.get(None), right.get(ctx), _BinNode.make_branch)
497                candidate(ctx, left.get(ctx), right.get(None), _BinNode.make_branch)
498            if not hole:
499                for ctx in sorted(set(ret) - set([None])):
500                    candidate(None, ctx, ret[ctx], _BinNode.make_default)