From 114177e7d8e1f8017baca397f943e11644b14d63 Mon Sep 17 00:00:00 2001
From: David Stern <vorporeal@gmail.com>
Date: Mon, 3 Jun 2024 14:11:35 -0400
Subject: [PATCH] Improve performance of `difflib`. (#1)

* Improve performance of difflib, primarily by avoiding heap allocations.

* Add doc comment to real_quick_ratio.
---
 Cargo.toml             |  4 ++++
 src/differ.rs          |  5 +++--
 src/sequencematcher.rs | 37 ++++++++++++++++++++++++++++++++-----
 tests/tests.rs         |  2 +-
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2edf6df..f99df9b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,8 +12,12 @@ include = [
     "**/*.rs",
     "Cargo.toml",
 ]
+edition = "2021"
 
 
 [[test]]
 name = "tests"
 
+[dependencies]
+rustc-hash = "1.1.0"
+
diff --git a/src/differ.rs b/src/differ.rs
index 5caef9a..166421a 100644
--- a/src/differ.rs
+++ b/src/differ.rs
@@ -1,6 +1,7 @@
-use sequencematcher::SequenceMatcher;
 use std::cmp;
-use utils::{count_leading, str_with_similar_chars};
+
+use crate::sequencematcher::SequenceMatcher;
+use crate::utils::{count_leading, str_with_similar_chars};
 
 #[derive(Default)]
 pub struct Differ {
diff --git a/src/sequencematcher.rs b/src/sequencematcher.rs
index 4157808..1bf77c0 100644
--- a/src/sequencematcher.rs
+++ b/src/sequencematcher.rs
@@ -1,7 +1,10 @@
 use std::cmp::{max, min};
 use std::collections::HashMap;
 use std::hash::Hash;
-use utils::calculate_ratio;
+
+use rustc_hash::FxHashMap;
+
+use crate::utils::calculate_ratio;
 
 #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
 pub struct Match {
@@ -57,6 +60,12 @@ pub struct SequenceMatcher<'a, T: 'a + Sequence> {
     opcodes: Option<Vec<Opcode>>,
     is_junk: Option<fn(&T) -> bool>,
     second_sequence_elements: HashMap<&'a T, Vec<usize>>,
+
+    // These hash maps are used by `find_longest_match`; we store them here
+    // to avoid making many unnecessary heap allocations; that function
+    // gets called many times as part of `get_close_matches`.
+    longest_match_j2len: FxHashMap<usize, usize>,
+    longest_match_new_j2len: FxHashMap<usize, usize>,
 }
 
 impl<'a, T: Sequence> SequenceMatcher<'a, T> {
@@ -71,6 +80,8 @@ impl<'a, T: Sequence> SequenceMatcher<'a, T> {
             opcodes: None,
             is_junk: None,
             second_sequence_elements: HashMap::new(),
+            longest_match_j2len: FxHashMap::default(),
+            longest_match_new_j2len: FxHashMap::default(),
         };
         matcher.set_seqs(first_sequence, second_sequence);
         matcher
@@ -138,7 +149,7 @@ impl<'a, T: Sequence> SequenceMatcher<'a, T> {
     }
 
     pub fn find_longest_match(
-        &self,
+        &mut self,
         first_start: usize,
         first_end: usize,
         second_start: usize,
@@ -148,14 +159,19 @@ impl<'a, T: Sequence> SequenceMatcher<'a, T> {
         let second_sequence = &self.second_sequence;
         let second_sequence_elements = &self.second_sequence_elements;
         let (mut best_i, mut best_j, mut best_size) = (first_start, second_start, 0);
-        let mut j2len: HashMap<usize, usize> = HashMap::new();
+        let mut j2len = &mut self.longest_match_j2len;
+        let mut new_j2len = &mut self.longest_match_new_j2len;
+
+        // Clear out any old data in our reusable hash map allocations.
+        j2len.clear();
+        new_j2len.clear();
+
         for (i, item) in first_sequence
             .iter()
             .enumerate()
             .take(first_end)
             .skip(first_start)
         {
-            let mut new_j2len: HashMap<usize, usize> = HashMap::new();
             if let Some(indexes) = second_sequence_elements.get(item) {
                 for j in indexes {
                     let j = *j;
@@ -180,7 +196,8 @@ impl<'a, T: Sequence> SequenceMatcher<'a, T> {
                     }
                 }
             }
-            j2len = new_j2len;
+            std::mem::swap(&mut j2len, &mut new_j2len);
+            new_j2len.clear();
         }
         for _ in 0..2 {
             while best_i > first_start
@@ -343,4 +360,14 @@ impl<'a, T: Sequence> SequenceMatcher<'a, T> {
             self.first_sequence.len() + self.second_sequence.len(),
         )
     }
+
+    /// Return an upper bound on ratio() very quickly.
+    ///
+    /// This isn't defined beyond that it is an upper bound on .ratio(), and
+    /// is faster to compute than either .ratio() or .quick_ratio().
+    pub fn real_quick_ratio(&self) -> f32 {
+        let la = self.first_sequence.len();
+        let lb = self.second_sequence.len();
+        calculate_ratio(min(la, lb), la + lb)
+    }
 }
diff --git a/tests/tests.rs b/tests/tests.rs
index 8f49de9..e85e341 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -5,7 +5,7 @@ use difflib::sequencematcher::{Match, Opcode, SequenceMatcher};
 
 #[test]
 fn test_longest_match() {
-    let matcher = SequenceMatcher::new(" abcd", "abcd abcd");
+    let mut matcher = SequenceMatcher::new(" abcd", "abcd abcd");
     let m = matcher.find_longest_match(0, 5, 0, 9);
     assert_eq!(m.first_start, 0);
     assert_eq!(m.second_start, 4);