Make CPDF_PageLabel::GetLabel() faster

Implement CPDF_NumberTree::GetLowerBound(), so that
CPDF_PageLabel::GetLabel() can find the value it needs in 1 search of
the number tree, instead of N. With this speed up, enable
PageLabelTest.GetLabelPerf, which now runs in ~1 second instead of ~200.

Since CPDF_PageLabel is the implementation for FPDF_GetPageLabel(), the
public API is faster too.

Bug: 347268873
Change-Id: Ie64d94be78227b975db897411c386de1babf7065
Reviewed-on: https://pdfium-review.googlesource.com/c/pdfium/+/120533
Reviewed-by: Tom Sepez <tsepez@chromium.org>
Commit-Queue: Lei Zhang <thestig@chromium.org>
Reviewed-by: Thomas Sepez <tsepez@google.com>
diff --git a/core/fpdfdoc/cpdf_numbertree.cpp b/core/fpdfdoc/cpdf_numbertree.cpp
index 5959782..5abfaed 100644
--- a/core/fpdfdoc/cpdf_numbertree.cpp
+++ b/core/fpdfdoc/cpdf_numbertree.cpp
@@ -6,6 +6,7 @@
 
 #include "core/fpdfdoc/cpdf_numbertree.h"
 
+#include <optional>
 #include <utility>
 
 #include "core/fpdfapi/parser/cpdf_array.h"
@@ -53,6 +54,55 @@
   return nullptr;
 }
 
+std::optional<CPDF_NumberTree::KeyValue> FindLowerBound(
+    const CPDF_Dictionary* node_dict,
+    int num) {
+  RetainPtr<const CPDF_Array> limits_array = node_dict->GetArrayFor("Limits");
+  if (limits_array) {
+    if (num < limits_array->GetIntegerAt(0)) {
+      return std::nullopt;
+    }
+    const int max_value = limits_array->GetIntegerAt(1);
+    if (num >= max_value) {
+      return CPDF_NumberTree::KeyValue(max_value,
+                                       FindNumberNode(node_dict, max_value));
+    }
+  }
+
+  RetainPtr<const CPDF_Array> numbers_array = node_dict->GetArrayFor("Nums");
+  if (numbers_array) {
+    for (size_t i = numbers_array->size() / 2; i > 0; --i) {
+      const size_t key_index = (i - 1) * 2;
+      const int key = numbers_array->GetIntegerAt(key_index);
+      if (num >= key) {
+        const size_t value_index = key_index + 1;
+        return CPDF_NumberTree::KeyValue(
+            key, numbers_array->GetDirectObjectAt(value_index));
+      }
+    }
+    return std::nullopt;
+  }
+
+  RetainPtr<const CPDF_Array> kids_array = node_dict->GetArrayFor("Kids");
+  if (!kids_array) {
+    return std::nullopt;
+  }
+
+  for (size_t i = kids_array->size(); i > 0; --i) {
+    RetainPtr<const CPDF_Dictionary> kid_dict = kids_array->GetDictAt(i - 1);
+    if (!kid_dict) {
+      continue;
+    }
+
+    std::optional<CPDF_NumberTree::KeyValue> result =
+        FindLowerBound(kid_dict.Get(), num);
+    if (result.has_value()) {
+      return result;
+    }
+  }
+  return std::nullopt;
+}
+
 }  // namespace
 
 CPDF_NumberTree::CPDF_NumberTree(RetainPtr<const CPDF_Dictionary> root)
@@ -63,3 +113,19 @@
 RetainPtr<const CPDF_Object> CPDF_NumberTree::LookupValue(int num) const {
   return FindNumberNode(root_.Get(), num);
 }
+
+std::optional<CPDF_NumberTree::KeyValue> CPDF_NumberTree::GetLowerBound(
+    int num) const {
+  return FindLowerBound(root_.Get(), num);
+}
+
+CPDF_NumberTree::KeyValue::KeyValue(int key, RetainPtr<const CPDF_Object> value)
+    : key(key), value(std::move(value)) {}
+
+CPDF_NumberTree::KeyValue::KeyValue(CPDF_NumberTree::KeyValue&&) noexcept =
+    default;
+
+CPDF_NumberTree::KeyValue& CPDF_NumberTree::KeyValue::operator=(
+    CPDF_NumberTree::KeyValue&&) noexcept = default;
+
+CPDF_NumberTree::KeyValue::~KeyValue() = default;
diff --git a/core/fpdfdoc/cpdf_numbertree.h b/core/fpdfdoc/cpdf_numbertree.h
index bcb712e..af20907 100644
--- a/core/fpdfdoc/cpdf_numbertree.h
+++ b/core/fpdfdoc/cpdf_numbertree.h
@@ -7,18 +7,43 @@
 #ifndef CORE_FPDFDOC_CPDF_NUMBERTREE_H_
 #define CORE_FPDFDOC_CPDF_NUMBERTREE_H_
 
+#include <optional>
+
 #include "core/fxcrt/retain_ptr.h"
 
 class CPDF_Dictionary;
 class CPDF_Object;
 
+// Represents a number tree that allows for sub-linear lookups of tree nodes.
+// See ISO 32000-1:2008 spec, section 7.9.7.
 class CPDF_NumberTree {
  public:
+  struct KeyValue {
+    KeyValue(int key, RetainPtr<const CPDF_Object> value);
+    KeyValue(KeyValue&) = delete;
+    KeyValue& operator=(KeyValue&) = delete;
+    KeyValue(KeyValue&&) noexcept;
+    KeyValue& operator=(KeyValue&&) noexcept;
+    ~KeyValue();
+
+    int key;
+    RetainPtr<const CPDF_Object> value;
+  };
+
   explicit CPDF_NumberTree(RetainPtr<const CPDF_Dictionary> root);
   ~CPDF_NumberTree();
 
+  // Finds the object in the number tree whose key is `num`. Returns nullptr in
+  // there is no `num` key in the tree.
   RetainPtr<const CPDF_Object> LookupValue(int num) const;
 
+  // Finds the object in the number tree with the largest key, such that
+  // `num` >= key. Returns the key/value pair if such a key exists, or
+  // std::nullopt otherwise.
+  // Note that this is similar to, but not exactly the same as
+  // std::lower_bound().
+  std::optional<KeyValue> GetLowerBound(int num) const;
+
  protected:
   RetainPtr<const CPDF_Dictionary> const root_;
 };
diff --git a/core/fpdfdoc/cpdf_pagelabel.cpp b/core/fpdfdoc/cpdf_pagelabel.cpp
index 29aa940..eb3cdd5 100644
--- a/core/fpdfdoc/cpdf_pagelabel.cpp
+++ b/core/fpdfdoc/cpdf_pagelabel.cpp
@@ -114,13 +114,10 @@
 
   CPDF_NumberTree number_tree(std::move(labels_dict));
   RetainPtr<const CPDF_Object> label_value;
-  int n = page_index;
-  while (n >= 0) {
-    label_value = number_tree.LookupValue(n);
-    if (label_value) {
-      break;
-    }
-    n--;
+  std::optional<CPDF_NumberTree::KeyValue> lower_bound =
+      number_tree.GetLowerBound(page_index);
+  if (lower_bound.has_value()) {
+    label_value = lower_bound.value().value;
   }
 
   const CPDF_Dictionary* label_dict =
@@ -135,7 +132,8 @@
   }
 
   ByteString style = label_dict->GetByteStringFor("S", ByteString());
-  int label_number = page_index - n + label_dict->GetIntegerFor("St", 1);
+  int label_number =
+      page_index - lower_bound.value().key + label_dict->GetIntegerFor("St", 1);
   label += GetLabelNumPortion(label_number, style);
   return label;
 }
diff --git a/core/fpdfdoc/cpdf_pagelabel_unittest.cpp b/core/fpdfdoc/cpdf_pagelabel_unittest.cpp
index 5a432b7..14aa513 100644
--- a/core/fpdfdoc/cpdf_pagelabel_unittest.cpp
+++ b/core/fpdfdoc/cpdf_pagelabel_unittest.cpp
@@ -201,8 +201,7 @@
   EXPECT_THAT(page_label()->GetLabel(10001), Eq(std::nullopt));
 }
 
-// TODO(crbug.com/347268873): Enable this test once GetLabel() is fast enough.
-TEST_F(PageLabelTest, DISABLED_GetLabelPerf) {
+TEST_F(PageLabelTest, GetLabelPerf) {
   for (int i = 0; i < 10001; ++i) {
     page_label()->GetLabel(i);
   }