3 files changed, 66 insertions, 61 deletions

diff --git a/match.c b/match.c
index 2a40472..da0f787 100644
--- a/match.c
+++ b/match.c
@@ -788,4 +788,50 @@ bool next_match(match_t **m, def_t *defs, file_t *f, pat_t *pat, pat_t *skip, bo
     return *m != NULL;
 }
 
+//
+// Helper function to track state while doing a depth-first search.
+//
+__attribute__((nonnull))
+static match_t *_get_numbered_capture(match_t *m, int *n)
+{
+    if (*n == 0) return m;
+    if (m->pat->type == BP_CAPTURE && m->pat->args.capture.namelen == 0) {
+        if (*n == 1) return m;
+        --(*n);
+    }
+    if (m->children) {
+        for (int i = 0; m->children[i]; i++) {
+            match_t *cap = _get_numbered_capture(m->children[i], n);
+            if (cap) return cap;
+        }
+    }
+    return NULL;
+}
+
+//
+// Get a specific numbered pattern capture.
+//
+match_t *get_numbered_capture(match_t *m, int n)
+{
+    return _get_numbered_capture(m, &n);
+}
+
+//
+// Get a capture with a specific name.
+//
+match_t *get_named_capture(match_t *m, const char *name, size_t namelen)
+{
+    if (m->pat->type == BP_CAPTURE && m->pat->args.capture.name
+        && m->pat->args.capture.namelen == namelen
+        && strncmp(m->pat->args.capture.name, name, m->pat->args.capture.namelen) == 0)
+        return m;
+    if (m->children) {
+        for (int i = 0; m->children[i]; i++) {
+            match_t *cap = get_named_capture(m->children[i], name, namelen);
+            if (cap) return cap;
+        }
+    }
+    return NULL;
+}
+
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/match.h b/match.h
index 535e20e..b6e7392 100644
--- a/match.h
+++ b/match.h
@@ -16,9 +16,12 @@ __attribute__((nonnull))
 void recycle_if_unused(match_t **at_m);
 size_t free_all_matches(void);
 size_t recycle_all_matches(void);
-
 bool next_match(match_t **m, def_t *defs, file_t *f, pat_t *pat, pat_t *skip, bool ignorecase);
 #define stop_matching(m) next_match(m, NULL, NULL, NULL, NULL, 0)
+__attribute__((nonnull))
+match_t *get_numbered_capture(match_t *m, int n);
+__attribute__((nonnull, pure))
+match_t *get_named_capture(match_t *m, const char *name, size_t namelen);
 
 #endif
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/print.c b/print.c
index a006a5e..1926b63 100644
--- a/print.c
+++ b/print.c
@@ -115,64 +115,6 @@ static const char *context_after(printer_t *pr, const char *pos)
 }
 
 //
-// Get a specific numbered pattern capture.
-//
-__attribute__((nonnull))
-static match_t *get_capture_by_num(match_t *m, int *n)
-{
-    if (*n == 0) return m;
-    if (m->pat->type == BP_CAPTURE && *n == 1) return m;
-    if (m->pat->type == BP_CAPTURE) --(*n);
-    if (m->children) {
-        for (int i = 0; m->children[i]; i++) {
-            match_t *cap = get_capture_by_num(m->children[i], n);
-            if (cap) return cap;
-        }
-    }
-    return NULL;
-}
-
-//
-// Get a capture with a specific name.
-//
-__attribute__((nonnull, pure))
-static match_t *get_capture_by_name(match_t *m, const char *name)
-{
-    if (m->pat->type == BP_CAPTURE && m->pat->args.capture.name
-        && strncmp(m->pat->args.capture.name, name, m->pat->args.capture.namelen) == 0)
-        return m;
-    if (m->children) {
-        for (int i = 0; m->children[i]; i++) {
-            match_t *cap = get_capture_by_name(m->children[i], name);
-            if (cap) return cap;
-        }
-    }
-    return NULL;
-}
-
-//
-// Get a capture by identifier (name or number).
-// Update *id to point to after the identifier (if found).
-//
-__attribute__((nonnull))
-static match_t *get_capture(match_t *m, const char **id)
-{
-    if (isdigit(**id)) {
-        int n = (int)strtol(*id, (char**)id, 10);
-        return get_capture_by_num(m, &n);
-    } else {
-        const char *end = after_name(*id);
-        if (end == *id) return NULL;
-        char *name = strndup(*id, (size_t)(end-*id));
-        match_t *cap = get_capture_by_name(m, name);
-        delete(&name);
-        *id = end;
-        if (**id == ';') ++(*id);
-        return cap;
-    }
-}
-
-//
 // Print the text of a match (no context).
 //
 static void _print_match(FILE *out, printer_t *pr, match_t *m)
@@ -194,9 +136,23 @@ static void _print_match(FILE *out, printer_t *pr, match_t *m)
             print_line_number(out, pr, line, pr->use_color ? color_replace : NULL, line > line_end);
 
             // Capture substitution
-            if (*r == '@' && r[1] && r[1] != '@') {
+            if (*r == '@' && r+1 < m->end && r[1] != '@') {
                 ++r;
-                match_t *cap = get_capture(m->children[0], &r);
+
+                // Retrieve the capture value:
+                match_t *cap = NULL;
+                if (isdigit(*r)) {
+                    int n = (int)strtol(r, (char**)&r, 10);
+                    cap = get_numbered_capture(m->children[0], n);
+                } else {
+                    const char *name = r, *end = after_name(r);
+                    if (end > name) {
+                        cap = get_named_capture(m->children[0], name, (size_t)(end - name));
+                        r = end;
+                        if (r < m->end && *r == ';') ++r;
+                    }
+                }
+
                 if (cap != NULL) {
                     _print_match(out, pr, cap);
                     if (pr->use_color && current_color != color_replace) {