Qualify function names (#678)

Closes #620

Author: tothtamas28

kmir/pyproject.toml

@@ -9,6 +9,7 @@ description = ""
 requires-python = "~=3.10"
 dependencies = [
     "kframework==v7.1.286",
+    "rust-demangler==1.0",
 ]
 
 [[project.authors]]

kmir/src/kmir/hello.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import logging
+from itertools import chain
 from math import ceil, log10
 from typing import TYPE_CHECKING
 
 from .smir import SMIRInfo
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
     from typing import Final
 
 _LOGGER: Final = logging.getLogger(__name__)
@@ -20,6 +22,8 @@ def link(smirs: list[SMIRInfo]) -> SMIRInfo:
     _LOGGER.info(f'Maximum type ID (offset) is {offset}, linking {len(smirs)} smir.json files')
 
     for i, smir in enumerate(smirs):
+        qualify_items(smir)
+
         smir_offset = offset * i
         _LOGGER.debug(f'Offset {smir_offset} for smir {smir._smir["name"]}')
         apply_offset(smir, smir_offset)
@@ -42,6 +46,135 @@ def id_range(smir: SMIRInfo) -> int:
     return max(0, *smir.function_symbols, *smir.types, *smir.spans, *smir.allocs)
 
 
+def qualify_items(info: SMIRInfo) -> None:
+    """Qualify each unqualified function item name.
+
+    The missing prefix is extracted from the symbol name.
+    """
+
+    for item in info._smir['items']:
+        match item:
+            case {
+                'symbol_name': symbol_name,
+                'mono_item_kind': {
+                    'MonoItemFn': {
+                        'name': name,
+                    } as mono_item_fn,
+                },
+            }:
+                qualified_name = _mono_item_fn_name(symbol_name=symbol_name, name=name)
+                if qualified_name != name:
+                    _LOGGER.info(f'Qualified item {symbol_name!r}: {name} -> {qualified_name}')
+                    mono_item_fn['name'] = qualified_name
+
+
+def _mono_item_fn_name(symbol_name: str, name: str) -> str:
+    """Extend ``name`` with a prefix from ``symbol_name``.
+
+    Example:
+        Symbol: foo :: bar :: do_something :: h0123456789abcdef
+        Name:          baz :: do_something :: <&u128>
+        Result: foo :: baz :: do_something :: <&u128>
+                ^^^    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+                |      |
+                |      +- kept from name
+                +- taken from symbol
+    """
+
+    def extract_id(s: str) -> str | None:
+        """Extract a Rust id prefix from a string."""
+        import re
+
+        m = re.match(r'^(?P<func>[a-zA-Z_][a-zA-Z0-9_]*)', s)
+        if not m:
+            return None
+        return m['func']
+
+    symbol = _demangle(symbol_name)
+    split_symbol = list(_symbol_segments(symbol))
+    split_name = list(_symbol_segments(name))
+
+    assert len(split_symbol) >= 2, 'The symbol name should contain at least two segments, an identifier and a hash'
+    # Extract the function name from `symbol_name`.
+    # It's the last segment with a valid id as prefix that's not the hash
+    i, fn_name = next(
+        ((i, fn_name) for i, s in enumerate(reversed(split_symbol[:-1])) if (fn_name := extract_id(s))), (None, None)
+    )
+    assert i is not None
+    assert fn_name is not None
+    symbol_index = len(split_symbol) - i - 2
+
+    # Find the index of the function name segment in the `split_name`
+    name_index = next((len(split_name) - i - 1 for i, s in enumerate(reversed(split_name)) if s == fn_name), None)
+    assert name_index is not None
+
+    if symbol_index < name_index:
+        # Do not add a prefix if the name prefix is longer than the symbol prefix
+        return name
+
+    # Construct the prefix and the result
+    return '::'.join(chain(split_symbol[: symbol_index - name_index], split_name))
+
+
+def _demangle(symbol: str) -> str:
+    import re
+
+    from rust_demangler import demangle  # type: ignore [import-untyped]
+
+    res = demangle(symbol)
+    res = re.sub(r'(?<!^)(?<!:)<', r'::<', res)  # insert '::' before '<' if not at the beginning or preceded by ':'
+    return res
+
+
+def _symbol_segments(s: str) -> Iterator[str]:
+    """Split a symbol at ``'::'`` not between ``'<'`` and ``'>'``."""
+    it = iter(s)
+    la = ''
+    buf: list[str] = []
+
+    def consume() -> None:
+        nonlocal la
+        la = next(it, '')
+
+    depth = 0
+    consume()
+    while la:
+        match la:
+            case ':':
+                consume()
+                match la:
+                    case ':':
+                        consume()
+                        if depth:
+                            buf += [':', ':']
+                        else:
+                            yield ''.join(buf)
+                            buf.clear()
+                    case '':
+                        buf.append(':')
+                        break
+                    case _:
+                        buf += [':', la]
+                        consume()
+            case '<':
+                buf.append(la)
+                consume()
+                depth += 1
+            case '>':
+                buf.append(la)
+                consume()
+                depth -= 1
+            case '':
+                raise AssertionError('The outer loop should ensure this is unreachable')
+            case _:
+                buf.append(la)
+                consume()
+
+    if depth != 0:
+        raise ValueError(f'Unbalanced <> in symbol: {s}')
+    yield ''.join(buf)
+
+
 def apply_offset(info: SMIRInfo, offset: int) -> None:
     # mutates the dictionary inside the SMIRInfo
     # all fields containing a `Ty` are updated, adding the given offset

kmir/src/kmir/linker.py

@@ -5,7 +5,6 @@
 │  (225 steps)
 ├─ 3 (terminal)
 │   #EndProgram ~> .K
-│   function: main
 │
 ┊  constraint: true
 ┊  subst: ...

…ests/single-bin/a_module::twice.expected …bin/single_exe::a_module::twice.expectedkmir/src/tests/integration/data/crate-tests/single-bin/a_module::twice.expected renamed to kmir/src/tests/integration/data/crate-tests/single-bin/single_exe::a_module::twice.expected kmir/src/tests/integration/data/crate-tests/single-bin/a_module::twice.expected renamed to kmir/src/tests/integration/data/crate-tests/single-bin/single_exe::a_module::twice.expected

@@ -5,7 +5,6 @@
 │  (730 steps)
 ├─ 3 (terminal)
 │   #EndProgram ~> .K
-│   function: main
 │
 ┊  constraint: true
 ┊  subst: ...

…ata/crate-tests/single-bin/main.expected …sts/single-bin/single_exe::main.expectedkmir/src/tests/integration/data/crate-tests/single-bin/main.expected renamed to kmir/src/tests/integration/data/crate-tests/single-bin/single_exe::main.expected kmir/src/tests/integration/data/crate-tests/single-bin/main.expected renamed to kmir/src/tests/integration/data/crate-tests/single-bin/single_exe::main.expected

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple
+
+import pytest
+
+if TYPE_CHECKING:
+    from typing import Final
+
+
+class _TestData(NamedTuple):
+    symbol: str
+    name: str
+    expected_demangled: str
+    expected_qualified: str
+
+
+TEST_DATA: Final[tuple[_TestData, ...]] = (
+    _TestData(
+        symbol='_ZN3foo3bar12do_something17h0123456789abcdefE',
+        name='baz::do_something::<&u128>',
+        expected_demangled='foo::bar::do_something::h0123456789abcdef',
+        expected_qualified='foo::baz::do_something::<&u128>',
+    ),
+    _TestData(
+        symbol='_ZN14small_test_lib3add17h67a713ff87afe55cE',
+        name='add',
+        expected_demangled='small_test_lib::add::h67a713ff87afe55c',
+        expected_qualified='small_test_lib::add',
+    ),
+    _TestData(
+        symbol='_ZN14small_test_lib6assume17h9c6e78f80476cd7fE',
+        name='assume',
+        expected_demangled='small_test_lib::assume::h9c6e78f80476cd7f',
+        expected_qualified='small_test_lib::assume',
+    ),
+    _TestData(
+        symbol='_ZN14small_test_lib7testing17test_add_in_range17h810360fe730dd322E',
+        name='testing::test_add_in_range',
+        expected_demangled='small_test_lib::testing::test_add_in_range::h810360fe730dd322',
+        expected_qualified='small_test_lib::testing::test_add_in_range',
+    ),
+    _TestData(
+        symbol='_ZN42_$LT$$RF$T$u20$as$u20$core..fmt..Debug$GT$3fmt17h3e83eb3d49c5b5a8E',
+        name='<&u128 as std::fmt::Debug>::fmt',
+        expected_demangled='<&T as core::fmt::Debug>::fmt::h3e83eb3d49c5b5a8',
+        expected_qualified='<&u128 as std::fmt::Debug>::fmt',
+    ),
+    _TestData(
+        symbol='_ZN4core3fmt3num51_$LT$impl$u20$core..fmt..Debug$u20$for$u20$u128$GT$3fmt17hfb9d82c7f008e36bE',
+        name='core::fmt::num::<impl std::fmt::Debug for u128>::fmt',
+        expected_demangled='core::fmt::num::<impl core::fmt::Debug for u128>::fmt::hfb9d82c7f008e36b',
+        expected_qualified='core::fmt::num::<impl std::fmt::Debug for u128>::fmt',
+    ),
+    _TestData(
+        symbol='_ZN4core3ptr29drop_in_place$LT$$RF$u128$GT$17hb88ca2b7d29bccccE',
+        name='std::ptr::drop_in_place::<&u128>',
+        expected_demangled='core::ptr::drop_in_place::<&u128>::hb88ca2b7d29bcccc',
+        expected_qualified='std::ptr::drop_in_place::<&u128>',
+    ),
+    _TestData(
+        symbol='_ZN4core9panicking13assert_failed17hdf64d315df90cf99E',
+        name='core::panicking::assert_failed::<u128, u128>',
+        expected_demangled='core::panicking::assert_failed::hdf64d315df90cf99',
+        expected_qualified='core::panicking::assert_failed::<u128, u128>',
+    ),
+    _TestData(
+        symbol='_ZN3std2rt10lang_start28_$u7b$$u7b$closure$u7d$$u7d$17h1f955f4ce95634f1E',
+        name='std::rt::lang_start::<()>::{closure#0}',
+        expected_demangled='std::rt::lang_start::{{closure}}::h1f955f4ce95634f1',
+        expected_qualified='std::rt::lang_start::<()>::{closure#0}',
+    ),
+    _TestData(
+        symbol='_ZN4core3ops8function6FnOnce40call_once$u7b$$u7b$vtable.shim$u7d$$u7d$17ha8802a5f1fccb12bE',
+        name='<{closure@std::rt::lang_start<()>::{closure#0}} as std::ops::FnOnce<()>>::call_once',
+        expected_demangled='core::ops::function::FnOnce::call_once{{vtable.shim}}::ha8802a5f1fccb12b',
+        expected_qualified='core::ops::function::<{closure@std::rt::lang_start<()>::{closure#0}} as std::ops::FnOnce<()>>::call_once',
+    ),
+    _TestData(
+        symbol='_ZN4core6option15Option$LT$T$GT$3map17h5b7d515e7f40e720E',
+        name='std::option::Option::<isize>::map::<crate1:MMyEnum, {closure@src/main.rs:19:29: 19:32}>',
+        expected_demangled='core::option::Option::<T>::map::h5b7d515e7f40e720',
+        expected_qualified='std::option::Option::<isize>::map::<crate1:MMyEnum, {closure@src/main.rs:19:29: 19:32}>',
+    ),
+)
+
+
+@pytest.mark.parametrize(
+    'symbol,_name,expected,_qualified',
+    TEST_DATA,
+    ids=[symbol for symbol, *_ in TEST_DATA],
+)
+def test_demangle(
+    symbol: str,
+    _name: str,
+    expected: str,
+    _qualified: str,
+) -> None:
+    from kmir.linker import _demangle
+
+    # When
+    actual = _demangle(symbol=symbol)
+
+    # Then
+    assert expected == actual
+
+
+@pytest.mark.parametrize(
+    'symbol_name,name,_demangled,expected',
+    TEST_DATA,
+    ids=[symbol_name for symbol_name, *_ in TEST_DATA],
+)
+def test_mono_item_fn_name(
+    symbol_name: str,
+    name: str,
+    _demangled: str,
+    expected: str,
+) -> None:
+    from kmir.linker import _mono_item_fn_name
+
+    # When
+    actual = _mono_item_fn_name(symbol_name=symbol_name, name=name)
+
+    # Then
+    assert expected == actual