55 lines
2 KiB
Python
55 lines
2 KiB
Python
|
"""
|
||
|
|
In-memory bidirectional message map for tracking relayed messages.
|
||
|
|
|
||
|
|
Maps original message IDs ↔ relayed copies across channels within a network.
|
||
|
|
Capped at MAP_LIMIT entries per direction to bound memory usage.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections import OrderedDict
|
||
|
|
from typing import Dict, List, Optional
|
||
|
|
|
||
|
|
from .config import MAP_LIMIT
|
||
|
|
|
||
|
|
|
||
|
|
class MessageMap:
|
||
|
|
"""Bidirectional map: original_id ↔ {channel_id: relayed_msg_id}.
|
||
|
|
|
||
|
|
``forward[net][original_id] = {ch_id: relayed_id, ...}``
|
||
|
|
``reverse[net][relayed_id] = original_id``
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.forward: Dict[str, OrderedDict[int, Dict[int, int]]] = {}
|
||
|
|
self.reverse: Dict[str, OrderedDict[int, int]] = {}
|
||
|
|
|
||
|
|
def add(self, network: str, original_id: int, mapping: Dict[int, int]) -> None:
|
||
|
|
"""Record that *original_id* was relayed to *mapping* channels."""
|
||
|
|
fwd = self.forward.setdefault(network, OrderedDict())
|
||
|
|
rev = self.reverse.setdefault(network, OrderedDict())
|
||
|
|
|
||
|
|
fwd[original_id] = mapping
|
||
|
|
for relayed_id in mapping.values():
|
||
|
|
rev[relayed_id] = original_id
|
||
|
|
|
||
|
|
# Prune oldest entries
|
||
|
|
while len(fwd) > MAP_LIMIT:
|
||
|
|
_old_orig, old_map = fwd.popitem(last=False)
|
||
|
|
for rid in old_map.values():
|
||
|
|
rev.pop(rid, None)
|
||
|
|
while len(rev) > MAP_LIMIT * 2:
|
||
|
|
rev.popitem(last=False)
|
||
|
|
|
||
|
|
def get_relayed(self, network: str, original_id: int) -> Dict[int, int]:
|
||
|
|
"""Get ``{ch_id: relayed_id}`` for an original message."""
|
||
|
|
return self.forward.get(network, {}).get(original_id, {})
|
||
|
|
|
||
|
|
def get_original(self, network: str, relayed_id: int) -> Optional[int]:
|
||
|
|
"""Get the original message ID from a relayed copy ID."""
|
||
|
|
return self.reverse.get(network, {}).get(relayed_id)
|
||
|
|
|
||
|
|
def get_all_relayed_ids(self, network: str, original_id: int) -> List[int]:
|
||
|
|
"""Get all relayed message IDs for an original."""
|
||
|
|
return list(self.get_relayed(network, original_id).values())
|