From 948af7b69348a3f921ced6c98d29f8bf5bb3bbea Mon Sep 17 00:00:00 2001
From: Stuart McCulloch <stuart.mcculloch@datadoghq.com>
Date: Sat, 17 Oct 2026 15:07:35 +0100
Subject: [PATCH] Preserve Unicode segments when extending trie leaves

Pass the leaf-segment context to appendLeaf instead of inferring it from
the next data character, which may belong to another branch and have the
leaf-marker bit set.

Cover characters around the marker boundary, short and longer
extensions, and subsequent insertions.
---
 .../instrument/utils/ClassNameTrie.java       |  7 +++--
 .../instrument/utils/ClassNameTrieTest.java   | 26 +++++++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java
index 3d205fa..dc1bccf 100644
--- a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java
+++ b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java
@@ -517,7 +517,7 @@ private void insertMapping(String key, char valueToInsert) {
 
         if ((value & LEAF_MARKER) != 0) {
           // append our new leaf node below the leaf branch and change it to a bud
-          jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert);
+          jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert, false);
           trieData[valueIndex] = (char) ((value & ~LEAF_MARKER) | BUD_MARKER);
           break;
         } else if ((value & BUD_MARKER) != 0) {
@@ -569,7 +569,7 @@ private void insertMapping(String key, char valueToInsert) {
           if ((value & LEAF_MARKER) != 0) {
             if (keyIndex < keyLength) {
               // key goes past leaf segment, extend old leaf with a bud node and add our new leaf
-              jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert);
+              jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert, true);
               trieData[valueIndex]--;
               break;
             } else {
@@ -825,12 +825,11 @@ private int insertLeafRight(
       return insertedCharacters;
     }
 
-    private int appendLeaf(int dataIndex, String key, int keyIndex, int value) {
+    private int appendLeaf(int dataIndex, String key, int keyIndex, int value, boolean insertBud) {
       int remainingKeyLength = key.length() - keyIndex;
       int insertedCharacters = 3 /* branch count, segment-length, value */ + remainingKeyLength;
 
       // if we're adding this to a leaf segment then we need to introduce a bud first
-      boolean insertBud = dataIndex < trieLength && (trieData[dataIndex] & LEAF_MARKER) != 0;
       if (insertBud) {
         insertedCharacters++;
       }
diff --git a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java
index ab86205..722202c 100644
--- a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java
+++ b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java
@@ -263,6 +263,32 @@ void tooManyLongJumps() {
         IllegalArgumentException.class, () -> builder.put("" + ((char) 67) + ((char) 2), 1));
   }
 
+  @ParameterizedTest
+  @ValueSource(chars = {'\u7FFF', '\u8000', '\uFFFF'})
+  void extendingLeafPreservesFollowingUnicodeSegment(char nextCharacter) {
+    for (String suffix : new String[] {"x", "xyz"}) {
+      ClassNameTrie.Builder builder = new ClassNameTrie.Builder();
+      String followingKey = "b" + nextCharacter;
+      builder.put("a", 1);
+      builder.put(followingKey, 2);
+
+      String extension = "a" + suffix;
+      builder.put(extension, 3);
+      assertEquals(1, builder.apply("a"));
+      assertEquals(2, builder.apply(followingKey));
+      assertEquals(3, builder.apply(extension));
+
+      builder.put(extension + "x", 4);
+      builder.put(followingKey + "x", 5);
+      ClassNameTrie trie = builder.buildTrie();
+      assertEquals(1, trie.apply("a"));
+      assertEquals(2, trie.apply(followingKey));
+      assertEquals(3, trie.apply(extension));
+      assertEquals(4, trie.apply(extension + "x"));
+      assertEquals(5, trie.apply(followingKey + "x"));
+    }
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void builtTriePreservesValuesAfterBuilderUpdates(boolean newBuilder) {