From 0269c7a12cb099a31f3d47e80f4582efc0fcfcac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20T=C3=B3th?=
 <tothtamas28@users.noreply.github.com>
Date: Fri, 20 Dec 2024 14:09:23 +0000
Subject: [PATCH] Generate an `abbrev` for each collection

---
 pyk/src/pyk/k2lean4/k2lean4.py | 27 ++++++++++++++++++++++++---
 pyk/src/pyk/kore/internal.py   |  8 +++++++-
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/pyk/src/pyk/k2lean4/k2lean4.py b/pyk/src/pyk/k2lean4/k2lean4.py
index d4b86146da..892aef8d7b 100644
--- a/pyk/src/pyk/k2lean4/k2lean4.py
+++ b/pyk/src/pyk/k2lean4/k2lean4.py
@@ -3,9 +3,10 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+from ..kore.internal import CollectionKind
 from ..kore.syntax import SortApp
 from ..utils import check_type
-from .model import Ctor, ExplBinder, Inductive, Module, Signature, Term
+from .model import Abbrev, Ctor, ExplBinder, Inductive, Module, Signature, Term
 
 if TYPE_CHECKING:
     from ..kore.internal import KoreDefn
@@ -19,7 +20,8 @@ class K2Lean4:
     def sort_module(self) -> Module:
         commands = []
         commands += self._inductives()
-        return Module(commands=self._inductives())
+        commands += self._collections()
+        return Module(commands=commands)
 
     def _inductives(self) -> list[Command]:
         def is_inductive(sort: str) -> bool:
@@ -41,7 +43,26 @@ def _inj_ctor(self, sort: str, subsort: str) -> Ctor:
         return Ctor(f'inj_{subsort}', Signature((ExplBinder(('x',), Term(subsort)),), Term(sort)))
 
     def _symbol_ctor(self, sort: str, symbol: str) -> Ctor:
-        param_sorts = (check_type(sort, SortApp).name for sort in self.defn.symbols[symbol].param_sorts)
+        param_sorts = (
+            check_type(sort, SortApp).name for sort in self.defn.symbols[symbol].param_sorts
+        )  # TODO eliminate check_type
         binders = tuple(ExplBinder((f'x{i}',), Term(sort)) for i, sort in enumerate(param_sorts))
         symbol = symbol.replace('-', '_')
         return Ctor(symbol, Signature(binders, Term(sort)))
+
+    def _collections(self) -> list[Command]:
+        return [self._collection(sort) for sort in sorted(self.defn.collections)]
+
+    def _collection(self, sort: str) -> Abbrev:
+        coll = self.defn.collections[sort]
+        elem = self.defn.symbols[coll.element]
+        sorts = ' '.join(check_type(sort, SortApp).name for sort in elem.param_sorts)  # TODO eliminate check_type
+        assert sorts
+        match coll.kind:
+            case CollectionKind.LIST:
+                val = Term(f'ListHook {sorts}')
+            case CollectionKind.MAP:
+                val = Term(f'MapHook {sorts}')
+            case CollectionKind.SET:
+                val = Term(f'SetHook {sorts}')
+        return Abbrev(sort, val, Signature((), Term('Type')))
diff --git a/pyk/src/pyk/kore/internal.py b/pyk/src/pyk/kore/internal.py
index b43d377c6d..95eb918bc7 100644
--- a/pyk/src/pyk/kore/internal.py
+++ b/pyk/src/pyk/kore/internal.py
@@ -189,7 +189,13 @@ def _config_symbols(self) -> set[str]:
             if sort in done:
                 continue
             done.add(sort)
-            symbols = self.constructors.get(sort, ())
+
+            symbols: list[str] = []
+            if sort in self.collections:
+                coll = self.collections[sort]
+                symbols += (coll.concat, coll.element, coll.unit)
+            symbols += self.constructors.get(sort, ())
+
             pending.extend(sort for symbol in symbols for sort in self._symbol_sorts(symbol))
             res.update(symbols)
         return res