Add RangeMap functionality to Pyk (#4482)

Depends on runtimeverification/llvm-backend#1096

Adds RangeMap functionality to Pyk analogous to all the existing
built-in collections

---------

Co-authored-by: Tamás Tóth <[email protected]>

pyk/src/pyk/kast/manip.py

@@ -84,7 +84,9 @@ def sort_assoc_label(label: str, kast: KInner) -> KInner:
 
 def sort_ac_collections(kast: KInner) -> KInner:
     def _sort_ac_collections(_kast: KInner) -> KInner:
-        if type(_kast) is KApply and (_kast.label.name in {'_Set_', '_Map_'} or _kast.label.name.endswith('CellMap_')):
+        if type(_kast) is KApply and (
+            _kast.label.name in {'_Set_', '_Map_', '_RangeMap_'} or _kast.label.name.endswith('CellMap_')
+        ):
             return sort_assoc_label(_kast.label.name, _kast)
         return _kast
 

pyk/src/pyk/kast/outer.py

@@ -859,7 +859,16 @@ def cell_collection_productions(self) -> tuple[KProduction, ...]:
     def _is_function(prod: KProduction) -> bool:
         def is_not_actually_function(label: str) -> bool:
             is_cell_map_constructor = label.endswith('CellMapItem') or label.endswith('CellMap_')
-            is_builtin_data_constructor = label in {'_Set_', '_List_', '_Map_', 'SetItem', 'ListItem', '_|->_'}
+            is_builtin_data_constructor = label in {
+                '_Set_',
+                '_List_',
+                '_Map_',
+                '_RangeMap_',
+                'SetItem',
+                'ListItem',
+                '_|->_',
+                '_r|->_',
+            }
             return is_cell_map_constructor or is_builtin_data_constructor
 
         return (Atts.FUNCTION in prod.att or Atts.FUNCTIONAL in prod.att) and not (

pyk/src/pyk/kore/match.py

@@ -86,6 +86,28 @@ def match_map(pattern: Pattern, *, cell: str | None = None) -> tuple[tuple[Patte
     return tuple(entries)
 
 
+def match_rangemap(pattern: Pattern) -> tuple[tuple[tuple[Pattern, Pattern], Pattern], ...]:
+    stop_symbol = "Lbl'Stop'RangeMap"
+    cons_symbol = "Lbl'Unds'RangeMap'Unds'"
+    item_symbol = "Lbl'Unds'r'Pipe'-'-GT-Unds'"
+
+    if type(pattern) is App:
+        match_app(pattern, stop_symbol)
+        return ()
+
+    assoc = match_left_assoc(pattern)
+    cons = match_app(assoc.app, cons_symbol)
+    items = (match_app(arg, item_symbol) for arg in cons.args)
+    entries = ((_match_range(item.args[0]), item.args[1]) for item in items)
+    return tuple(entries)
+
+
+def _match_range(pattern: Pattern) -> tuple[Pattern, Pattern]:
+    range_symbol = "LblRangeMap'Coln'Range"
+    range = match_app(pattern, range_symbol)
+    return (range.args[0], range.args[1])
+
+
 def kore_bool(pattern: Pattern) -> bool:
     dv = match_dv(pattern, BOOL)
     match dv.value.value:
@@ -253,6 +275,16 @@ def res(pattern: Pattern) -> tuple[tuple[K, V], ...]:
     return res
 
 
+def kore_rangemap_of(
+    key: Callable[[Pattern], K],
+    value: Callable[[Pattern], V],
+) -> Callable[[Pattern], tuple[tuple[tuple[K, K], V], ...]]:
+    def res(pattern: Pattern) -> tuple[tuple[tuple[K, K], V], ...]:
+        return tuple(((key(k[0]), key(k[1])), value(v)) for k, v in match_rangemap(pattern))
+
+    return res
+
+
 def case_symbol(
     *cases: tuple[str, Callable[[App], T]],
     default: Callable[[App], T] | None = None,

pyk/src/pyk/kore/prelude.py

@@ -252,6 +252,24 @@ def map_pattern(*args: tuple[Pattern, Pattern], cell: str | None = None) -> Patt
     return LeftAssoc(App(cons_symbol, args=(App(item_symbol, args=arg) for arg in args)))
 
 
+STOP_RANGEMAP: Final = App("Lbl'Stop'RangeMap")
+LBL_RANGEMAP: Final = SymbolId("Lbl'Unds'RangeMap'Unds'")
+LBL_RANGEMAP_ITEM: Final = SymbolId("Lbl'Unds'r'Pipe'-'-GT-Unds'")
+LBL_RANGEMAP_RANGE: Final = SymbolId("LblRangeMap'Coln'Range")
+
+
+def rangemap_pattern(*args: tuple[tuple[Pattern, Pattern], Pattern]) -> Pattern:
+    if not args:
+        return STOP_RANGEMAP
+
+    return LeftAssoc(
+        App(
+            LBL_RANGEMAP,
+            args=(App(LBL_RANGEMAP_ITEM, args=(App(LBL_RANGEMAP_RANGE, args=arg[0]), arg[1])) for arg in args),
+        )
+    )
+
+
 # ----
 # JSON
 # ----

pyk/src/pyk/prelude/collections.py

@@ -13,6 +13,7 @@
 SET: Final = KSort('Set')
 LIST: Final = KSort('List')
 MAP: Final = KSort('Map')
+RANGEMAP: Final = KSort('RangeMap')
 BAG: Final = KSort('Bag')
 
 
@@ -51,3 +52,16 @@ def map_item(k: KInner, v: KInner) -> KInner:
 def map_of(ks: dict[KInner, KInner] | Iterable[tuple[KInner, KInner]]) -> KInner:
     ks = dict(ks)
     return build_assoc(map_empty(), KLabel('_Map_'), (map_item(k, v) for k, v in ks.items()))
+
+
+def rangemap_empty() -> KInner:
+    return KApply('.RangeMap')
+
+
+def rangemap_item(k: tuple[KInner, KInner], v: KInner) -> KInner:
+    return KApply('_r|->_', (KApply('RangeMap:Range', k), v))
+
+
+def rangemap_of(ks: dict[tuple[KInner, KInner], KInner] | Iterable[tuple[tuple[KInner, KInner], KInner]]) -> KInner:
+    ks_dict: dict[tuple[KInner, KInner], KInner] = dict(ks)
+    return build_assoc(rangemap_empty(), KLabel('_RangeMap_'), (rangemap_item(k, v) for k, v in ks_dict.items()))

pyk/src/tests/unit/kore/test_match.py

@@ -15,10 +15,11 @@
     kore_int,
     kore_list_of,
     kore_map_of,
+    kore_rangemap_of,
     kore_set_of,
     kore_str,
 )
-from pyk.kore.prelude import dv, list_pattern, map_pattern, set_pattern
+from pyk.kore.prelude import dv, list_pattern, map_pattern, rangemap_pattern, set_pattern
 from pyk.kore.syntax import App
 
 if TYPE_CHECKING:
@@ -82,6 +83,12 @@ def res(*args: Pattern) -> App:
         kore_map_of(kore_int, kore_str, cell='ACell'),
         ((0, 'a'), (1, 'b'), (2, 'c')),
     ),
+    (rangemap_pattern(), kore_rangemap_of(kore_int, kore_str), ()),
+    (
+        rangemap_pattern(((dv(0), dv(1)), dv('a')), ((dv(2), dv(3)), dv('b'))),
+        kore_rangemap_of(kore_int, kore_str),
+        (((0, 1), 'a'), ((2, 3), 'b')),
+    ),
     (
         a(),
         case_symbol(