diff --git a/pyk/src/pyk/kast/manip.py b/pyk/src/pyk/kast/manip.py index fc453654d33..92c8ae03a3b 100644 --- a/pyk/src/pyk/kast/manip.py +++ b/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 diff --git a/pyk/src/pyk/kast/outer.py b/pyk/src/pyk/kast/outer.py index f6c6dfa7a07..9cfd4fc7e27 100644 --- a/pyk/src/pyk/kast/outer.py +++ b/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 ( diff --git a/pyk/src/pyk/kore/match.py b/pyk/src/pyk/kore/match.py index 86c18aeca44..4f6534e2d23 100644 --- a/pyk/src/pyk/kore/match.py +++ b/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_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 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 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, diff --git a/pyk/src/pyk/kore/prelude.py b/pyk/src/pyk/kore/prelude.py index 6b8bd629b2e..aec8c6ffd09 100644 --- a/pyk/src/pyk/kore/prelude.py +++ b/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 # ---- diff --git a/pyk/src/pyk/prelude/collections.py b/pyk/src/pyk/prelude/collections.py index b1d93adbc1b..bb6351145d8 100644 --- a/pyk/src/pyk/prelude/collections.py +++ b/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())) diff --git a/pyk/src/tests/unit/kore/test_match.py b/pyk/src/tests/unit/kore/test_match.py index 7822164d323..f59cc9aef79 100644 --- a/pyk/src/tests/unit/kore/test_match.py +++ b/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(