From ccfc7ea3bb5e5c57cc560f6d66202536ce54a140 Mon Sep 17 00:00:00 2001
From: Michael Baikov <manpacket@gmail.com>
Date: Thu, 10 Oct 2024 13:10:44 -0400
Subject: [PATCH] Make sure not to drop used labels

Regexp find returns the whole group, including non capturing part, so we
can't use it directly. This MR makes regexes more private and uses
capture groups to extract labels

Fixes #318
---
 Changelog.md    |  1 +
 src/asm.rs      |  8 ++------
 src/demangle.rs | 33 +++++++++++++++++++++++++--------
 3 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/Changelog.md b/Changelog.md
index 9620f75..bdc4543 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -1,6 +1,7 @@
 # Change Log
 
 ## [0.2.41] - Unreleased
+- make sure not to drop used labels (#318)
 - add release-lto profile for slightly smaller/faster version
   thanks @zamazan4ik for the suggestion
 - detect and render merged functions (#310)
diff --git a/src/asm.rs b/src/asm.rs
index eb5c5d8..1ca94d8 100644
--- a/src/asm.rs
+++ b/src/asm.rs
@@ -263,7 +263,6 @@ fn used_labels<'a>(stmts: &'_ [Statement<'a>]) -> BTreeSet<&'a str> {
             Statement::Dunno(s) => Some(s),
         })
         .flat_map(demangle::local_labels)
-        .map(|m| m.as_str())
         .collect::<BTreeSet<_>>()
 }
 
@@ -629,11 +628,8 @@ impl<'a> Dumpable for Asm<'a> {
                     })
                     | Statement::Directive(Directive::Generic(GenericDirective(arg))) = s
                     {
-                        for label in crate::demangle::local_labels_reg().find_iter(arg) {
-                            let referenced_label = label.as_str().trim();
-                            if let Some(constant_range) =
-                                scan_constant(referenced_label, &sections, lines)
-                            {
+                        for label in crate::demangle::local_labels(arg) {
+                            if let Some(constant_range) = scan_constant(label, &sections, lines) {
                                 if !seen.contains(&constant_range)
                                     && !print_range.fully_contains(constant_range)
                                 {
diff --git a/src/demangle.rs b/src/demangle.rs
index b6d8d27..0e44fe4 100644
--- a/src/demangle.rs
+++ b/src/demangle.rs
@@ -1,3 +1,6 @@
+#![allow(clippy::needless_pub_self)] // default is wrong, I want to mark some items as explicitly private
+                                     // otherwise pub(self) makes no sense
+
 use crate::{color, opts::NameDisplay};
 use owo_colors::OwoColorize;
 use regex::{Regex, RegexSet, Replacer};
@@ -20,7 +23,7 @@ pub fn demangled(input: &str) -> Option<Demangle> {
     Some(name)
 }
 
-const GLOBAL_LABELS_REGEX: &str = r"\b_?(_[a-zA-Z0-9_$\.]+)";
+pub(self) const GLOBAL_LABELS_REGEX: &str = r"\b_?(_[a-zA-Z0-9_$\.]+)";
 
 // This regex is two parts
 // 1. \.L[a-zA-Z0-9_$\.]+
@@ -35,22 +38,22 @@ const GLOBAL_LABELS_REGEX: &str = r"\b_?(_[a-zA-Z0-9_$\.]+)";
 // there  as long as it doesn't look like a label.
 //
 // Note: this rejects "labels" like `H.Lfoo` but accepts `.Lexception` and `[some + .Label]`
-const LOCAL_LABELS_REGEX: &str = r"(?:[^\w\d\$\.]|^)(\.L[a-zA-Z0-9_\$\.]+|\bLBB[0-9_]+)";
+pub(self) const LOCAL_LABELS_REGEX: &str = r"(?:[^\w\d\$\.]|^)(\.L[a-zA-Z0-9_\$\.]+|\bLBB[0-9_]+)";
 
 // temporary labels
-const TEMP_LABELS_REGEX: &str = r"\b(Ltmp[0-9]+)\b";
+pub(self) const TEMP_LABELS_REGEX: &str = r"\b(Ltmp[0-9]+)\b";
 
-fn global_labels_reg() -> &'static Regex {
+pub(self) fn global_labels_reg() -> &'static Regex {
     static GLOBAL_LABELS: OnceLock<Regex> = OnceLock::new();
     GLOBAL_LABELS.get_or_init(|| Regex::new(GLOBAL_LABELS_REGEX).expect("regexp should be valid"))
 }
 
-pub(crate) fn local_labels_reg() -> &'static Regex {
+pub(self) fn local_labels_reg() -> &'static Regex {
     static LOCAL_LABELS: OnceLock<Regex> = OnceLock::new();
     LOCAL_LABELS.get_or_init(|| Regex::new(LOCAL_LABELS_REGEX).expect("regexp should be valid"))
 }
 
-fn label_kinds_reg() -> &'static RegexSet {
+pub(self) fn label_kinds_reg() -> &'static RegexSet {
     static LABEL_KINDS: OnceLock<RegexSet> = OnceLock::new();
     LABEL_KINDS.get_or_init(|| {
         RegexSet::new([LOCAL_LABELS_REGEX, GLOBAL_LABELS_REGEX, TEMP_LABELS_REGEX])
@@ -66,8 +69,22 @@ pub enum LabelKind {
     Unknown,
 }
 
-pub fn local_labels(input: &str) -> regex::Matches {
-    local_labels_reg().find_iter(input)
+#[test]
+fn local_labels_works() {
+    let s0 = "vmovaps xmm0, xmmword ptr [rip + .LCPI0_0]";
+    assert_eq!(local_labels(s0).collect::<Vec<_>>(), [".LCPI0_0"]);
+
+    let s1 = "vmovaps xmm0, xmmword ptr [rip + B.LCPI0_0]";
+    assert_eq!(local_labels(s1).collect::<Vec<_>>(), [] as [&str; 0]);
+
+    let s2 = ".Lexception";
+    assert_eq!(local_labels(s2).collect::<Vec<_>>(), [".Lexception"]);
+}
+
+pub(crate) fn local_labels(input: &str) -> impl Iterator<Item = &str> {
+    local_labels_reg()
+        .captures_iter(input)
+        .filter_map(|c| Some(c.get(1)?.as_str()))
 }
 
 #[must_use]