Access violation fix

Author: apsonawane

operators/tokenizer/case_encoder.h

@@ -170,49 +170,71 @@ class CaseEncoder {
   void PostProcess(std::string* normalized, std::vector<size_t>* norm_to_orig) {
     if (!seen_three_spans_) return;
 
+    // Safety check: ensure norm_to_orig has enough elements
+    if (norm_to_orig->size() < normalized->size()) {
+      return;  // Cannot safely process - sizes don't match
+    }
+
     std::string normalized_temp;
     normalized_temp.reserve(normalized->size());
 
     std::vector<size_t> norm_to_orig_temp;
     norm_to_orig_temp.reserve(norm_to_orig->size());
 
     const char* sig_it = signature_.data();
+    const char* sig_end = signature_.data() + signature_.length();
 
     auto nrm_it = normalized->cbegin();
+    auto nrm_end = normalized->cend();
     auto n2o_it = norm_to_orig->cbegin();
+    auto n2o_end = norm_to_orig->cend();
 
     for (const auto& span : Search(signature_)) {
       size_t len = std::distance(sig_it, span.first);
 
+      // Bounds check before advancing iterators
+      if (std::distance(nrm_it, nrm_end) < static_cast<ptrdiff_t>(len) ||
+          std::distance(n2o_it, n2o_end) < static_cast<ptrdiff_t>(len)) {
+        break;  // Not enough elements remaining
+      }
+
       normalized_temp.insert(normalized_temp.end(), nrm_it, nrm_it + len);
       norm_to_orig_temp.insert(norm_to_orig_temp.end(), n2o_it, n2o_it + len);
 
       sig_it += len;
       nrm_it += len;
       n2o_it += len;
+
+      // Bounds check before dereferencing
+      if (n2o_it == n2o_end) break;
+
       normalized_temp.push_back(cAllUppercase);
       norm_to_orig_temp.push_back(*n2o_it);
 
       while (sig_it != span.second) {
+        if (sig_it >= sig_end || nrm_it >= nrm_end || n2o_it >= n2o_end) break;
+
         if (*sig_it == cUppercase) {
           sig_it++;
           nrm_it++;
           n2o_it++;
         }
+        if (sig_it >= sig_end || nrm_it >= nrm_end || n2o_it >= n2o_end) break;
+
         sig_it++;
         normalized_temp.push_back(*nrm_it++);
         norm_to_orig_temp.push_back(*n2o_it++);
       }
-      if (sig_it != signature_.data() + signature_.length()) {
-        if (*sig_it != cUppercase) {
+      if (sig_it != sig_end) {
+        if (*sig_it != cUppercase && n2o_it != n2o_end) {
           normalized_temp.push_back(cLowercase);
           norm_to_orig_temp.push_back(*n2o_it);
         }
       }
     }
 
-    if (nrm_it != normalized->cend()) normalized_temp.insert(normalized_temp.end(), nrm_it, normalized->cend());
-    if (n2o_it != norm_to_orig->cend()) norm_to_orig_temp.insert(norm_to_orig_temp.end(), n2o_it, norm_to_orig->cend());
+    if (nrm_it != nrm_end) normalized_temp.insert(normalized_temp.end(), nrm_it, nrm_end);
+    if (n2o_it != n2o_end) norm_to_orig_temp.insert(norm_to_orig_temp.end(), n2o_it, n2o_end);
 
     normalized->swap(normalized_temp);
     norm_to_orig->swap(norm_to_orig_temp);

operators/tokenizer/ugm_kernels.hpp

@@ -537,7 +537,9 @@ struct SpmUgmTokenizer {
   std::string NmtNormalize(const std::string& input) const {
     std::string normalized;
     normalized.reserve(input.size() * 3);
-    std::vector<size_t> norm_to_orig(input.size() * 3);
+    // Use a vector that tracks original positions - reserve capacity but start empty
+    std::vector<size_t> norm_to_orig;
+    norm_to_orig.reserve(input.size() * 3);
 
     const std::string space = tokenizer_escape_whitespaces_ ? std::string(spm_escaped_space) : " ";
 
@@ -551,40 +553,60 @@ struct SpmUgmTokenizer {
     size_t input_len = input.size();
 
     std::string_view input_view(input);
-    int consumed = 0;
+    size_t orig_offset = 0;
 
     while (!input_view.empty()) {
       auto p = case_encoder_->NormalizePrefix(input_view);
 
+      // Safety check: if nothing was consumed and nothing was returned, skip one byte to avoid infinite loop
+      if (p.second == 0 && p.first.empty()) {
+        // Advance by one UTF-8 character to prevent infinite loop
+        size_t skip = std::min(ustring::UTF8Len(input_view[0]), input_view.size());
+        orig_offset += skip;
+        input_view.remove_prefix(skip);
+        continue;
+      }
+
       for (size_t i = 0; i < p.first.size(); i++) {
         char c = p.first[i];
         if (c != ' ') {
           if (!processing_non_ws) {
             processing_non_ws = true;
             if ((shall_prepend_space && !is_space_prepended) || shall_merge_spaces) {
               normalized.append(space);
+              // Track original position for space characters
+              for (size_t j = 0; j < space.size(); j++) {
+                norm_to_orig.push_back(orig_offset);
+              }
               is_space_prepended = true;
             }
           }
           normalized.push_back(c);
+          norm_to_orig.push_back(orig_offset + i);
         } else {
           if (processing_non_ws) {
             processing_non_ws = false;
           }
           if (!shall_merge_spaces) {
             normalized.append(space);
+            for (size_t j = 0; j < space.size(); j++) {
+              norm_to_orig.push_back(orig_offset + i);
+            }
           }
         }
       }
 
-      consumed += p.second;
-      input_view.remove_prefix(p.second);
+      orig_offset += static_cast<size_t>(p.second);
+      input_view.remove_prefix(static_cast<size_t>(p.second));
     }
 
     case_encoder_->PostProcess(&normalized, &norm_to_orig);
 
     if (shall_append_space) {
       normalized.append(space);
+      for (size_t j = 0; j < space.size(); j++) {
+        norm_to_orig.push_back(input.size());
+      }
     }
 
     return normalized;