diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 823fe4f..d7ed13b 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -45,7 +45,7 @@ jobs: uses: PyO3/maturin-action@v1 with: target: ${{ matrix.platform.target }} - args: --release --out dist -i python3.9 python3.10 python3.11 python3.12 python3.13 + args: --release --out dist -i python3.9 python3.10 python3.11 python3.12 python3.13 python3.14 sccache: 'true' manylinux: auto - name: Upload wheels @@ -76,7 +76,7 @@ jobs: uses: PyO3/maturin-action@v1 with: target: ${{ matrix.platform.target }} - args: --release --out dist -i python3.9 python3.10 python3.11 python3.12 python3.13 + args: --release --out dist -i python3.9 python3.10 python3.11 python3.12 python3.13 python3.14 sccache: 'true' manylinux: musllinux_1_2 - name: Upload wheels @@ -116,6 +116,10 @@ jobs: with: python-version: '3.13' architecture: ${{ matrix.platform.target }} + - uses: actions/setup-python@v5 + with: + python-version: '3.14' + architecture: ${{ matrix.platform.target }} - name: Build wheels uses: PyO3/maturin-action@v1 env: @@ -156,6 +160,9 @@ jobs: - uses: actions/setup-python@v5 with: python-version: '3.13' + - uses: actions/setup-python@v5 + with: + python-version: '3.14' - name: Build wheels uses: PyO3/maturin-action@v1 env: diff --git a/Cargo.lock b/Cargo.lock index 44c5b5a..291a05a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,16 +1,25 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + [[package]] name = "anes" version = "0.1.6" @@ -19,15 +28,15 @@ checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" [[package]] name = "anstyle" -version = "1.0.8" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bec1de6f59aedf83baf9ff929c98f2ad654b97c9510f4e70cf6f661d49fd5b1" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" [[package]] name = "autocfg" -version = "1.3.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "bincode" @@ -40,7 +49,7 @@ dependencies = [ [[package]] name = "blitztext" -version = "0.1.6" +version = "0.2.0" dependencies = [ "bincode", "criterion", @@ -54,15 +63,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.16.0" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" - -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "cast" @@ -70,11 +73,32 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "cc" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877bf67ac8e52449eb2f9e98a344ea163beb70dd1f9eebe3cc5357dcfeaee4df" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core", +] [[package]] name = "ciborium" @@ -105,18 +129,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.15" +version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11d8838454fda655dafd3accb2b6e2bea645b9e4078abe84a22ceb947235c5cc" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.5.15" +version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "216aec2b177652e3846684cbfe25c9964d18ec45234f0f5da5157b207ed1aab6" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" dependencies = [ "anstyle", "clap_lex", @@ -124,31 +148,39 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.2" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + +[[package]] +name = "cpufeatures" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] [[package]] name = "criterion" -version = "0.5.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" dependencies = [ + "alloca", "anes", "cast", "ciborium", "clap", "criterion-plot", - "is-terminal", "itertools", "num-traits", - "once_cell", "oorandom", + "page_size", "plotters", "rayon", "regex", "serde", - "serde_derive", "serde_json", "tinytemplate", "walkdir", @@ -156,9 +188,9 @@ dependencies = [ [[package]] name = "criterion-plot" -version = "0.5.0" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" dependencies = [ "cast", "itertools", @@ -166,9 +198,9 @@ dependencies = [ [[package]] name = "crossbeam-deque" -version = "0.8.5" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a" dependencies = [ "crossbeam-epoch", "crossbeam-utils", @@ -176,51 +208,82 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d" dependencies = [ "crossbeam-utils", ] [[package]] name = "crossbeam-utils" -version = "0.8.20" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" [[package]] name = "crunchy" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "either" -version = "1.13.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "find-msvc-tools" +version = "0.1.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c66915cc5391097de7a2b0b972e8505aa33c0df401fd3a6811dbdb32881ba48" + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] [[package]] name = "getrandom" -version = "0.3.4" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", "libc", "r-efi", - "wasip2", + "rand_core", ] [[package]] name = "half" -version = "2.4.1" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", + "zerocopy", ] [[package]] @@ -229,79 +292,43 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "hermit-abi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" - -[[package]] -name = "indoc" -version = "2.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" - -[[package]] -name = "is-terminal" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f23ff5ef2b80d608d61efee834934d862cd92461afc0560dedf493e4c033738b" -dependencies = [ - "hermit-abi", - "libc", - "windows-sys 0.52.0", -] - [[package]] name = "itertools" -version = "0.10.5" +version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" dependencies = [ "either", ] [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.69" +version = "0.3.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29c15563dc2726973df627357ce0c9ddddbea194836909d655df6a75d2cf296d" +checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5" dependencies = [ + "cfg-if", + "futures-util", "wasm-bindgen", ] [[package]] name = "libc" -version = "0.2.155" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c" - -[[package]] -name = "log" -version = "0.4.22" +version = "0.2.191" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" [[package]] name = "memchr" -version = "2.7.4" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" - -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "num-traits" @@ -314,21 +341,37 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.19.0" +version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "oorandom" -version = "11.1.4" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b410bbe7e14ab526a0e86877eb47c6996a2bd7746f027ba551028c925390e4e9" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "plotters" -version = "0.3.6" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a15b6eccb8484002195a3e44fe65a4ce8e93a625797a063735536fd59cb01cf3" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" dependencies = [ "num-traits", "plotters-backend", @@ -339,76 +382,62 @@ dependencies = [ [[package]] name = "plotters-backend" -version = "0.3.6" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "414cec62c6634ae900ea1c56128dfe87cf63e7caece0852ec76aba307cebadb7" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" [[package]] name = "plotters-svg" -version = "0.3.6" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81b30686a7d9c3e010b84284bdd26a29f2138574f52f5eb6f794fc0ad924e705" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" dependencies = [ "plotters-backend", ] [[package]] name = "portable-atomic" -version = "1.7.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da544ee218f0d287a911e9c99a39a8c9bc8fcad3cb8db5959940044ecfc67265" - -[[package]] -name = "ppv-lite86" -version = "0.2.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" -dependencies = [ - "zerocopy", -] +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" [[package]] name = "proc-macro2" -version = "1.0.86" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] [[package]] name = "pyo3" -version = "0.22.6" +version = "0.29.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884" +checksum = "700d18fa267b73b9b521fd7e13580e2f446916f176cacee1ab63fcc8191f1655" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-build-config" -version = "0.22.6" +version = "0.29.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38" +checksum = "7b3fc0c4d08f6bb10e71fe39dfb9e2f59c6eb6854e22ec8092f50c69a4499adb" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.22.6" +version = "0.29.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636" +checksum = "dfc0b8e19df29aad7086cf977bb0c2a2f143e30567eb113e9cf72b62ca698330" dependencies = [ "libc", "pyo3-build-config", @@ -416,78 +445,65 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.22.6" +version = "0.29.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453" +checksum = "6100e8a4b5eba53afaa5ed078364851a0b2499c553a44c31026b929049b49dc6" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "pyo3-macros-backend" -version = "0.22.6" +version = "0.29.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe" +checksum = "6143877a16e82b5a727b7127ff4cd86858a24a28d745d72f43e6f227c7b1bdb3" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "quote" -version = "1.0.36" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fa76aaf39101c457836aec0ce2316dbdc3ab723cdda1c6bd4e6ad4208acaca7" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] [[package]] name = "r-efi" -version = "5.3.0" +version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" -dependencies = [ - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.9.0" +version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" dependencies = [ - "ppv-lite86", + "chacha20", + "getrandom", "rand_core", ] [[package]] name = "rand_core" -version = "0.9.3" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" -dependencies = [ - "getrandom", -] +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] name = "rayon" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" dependencies = [ "either", "rayon-core", @@ -505,9 +521,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.2" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -517,9 +533,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.13" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -528,21 +544,21 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.8" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "rustc-hash" -version = "2.1.1" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] -name = "ryu" -version = "1.0.18" +name = "rustversion" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "same-file" @@ -555,9 +571,9 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -565,41 +581,65 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.9", ] [[package]] name = "serde_json" -version = "1.0.124" +version = "1.0.154" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66ad62847a56b3dba58cc891acd13884b9c61138d330c0d7b6181713d4fce38d" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" dependencies = [ "itoa", "memchr", - "ryu", "serde", + "serde_core", + "zmij", ] +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "syn" -version = "2.0.87" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" dependencies = [ "proc-macro2", "quote", @@ -608,9 +648,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "tinytemplate" @@ -624,15 +664,9 @@ dependencies = [ [[package]] name = "unicode-ident" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" - -[[package]] -name = "unindent" -version = "0.2.3" +version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" [[package]] name = "walkdir" @@ -644,45 +678,24 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "wasip2" -version = "1.0.1+wasi-0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" -dependencies = [ - "wit-bindgen", -] - [[package]] name = "wasm-bindgen" -version = "0.2.92" +version = "0.2.129" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4be2531df63900aeb2bca0daaaddec08491ee64ceecbee5076636a3b026795a8" +checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409" dependencies = [ "cfg-if", - "wasm-bindgen-macro", -] - -[[package]] -name = "wasm-bindgen-backend" -version = "0.2.92" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "614d787b966d3989fa7bb98a654e369c762374fd3213d212cfc0251257e747da" -dependencies = [ - "bumpalo", - "log", "once_cell", - "proc-macro2", - "quote", - "syn", + "rustversion", + "wasm-bindgen-macro", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.92" +version = "0.2.129" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1f8823de937b71b9460c0c34e25f3da88250760bec0ebac694b49997550d726" +checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -690,147 +703,104 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.92" +version = "0.2.129" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7" +checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7" dependencies = [ + "bumpalo", "proc-macro2", "quote", - "syn", - "wasm-bindgen-backend", + "syn 3.0.9", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.92" +version = "0.2.129" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af190c94f2773fdb3729c55b007a722abb5384da03bc0986df4c289bf5567e96" +checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6" +dependencies = [ + "unicode-ident", +] [[package]] name = "web-sys" -version = "0.3.69" +version = "0.3.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77afa9a11836342370f4817622a2f0f418b134426d91a82dfb48f532d2ec13ef" +checksum = "88261b9deccee56594c11a3460c462c41f58d148598fe70ad77070126a68aba4" dependencies = [ "js-sys", "wasm-bindgen", ] [[package]] -name = "winapi-util" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" -dependencies = [ - "windows-sys 0.59.0", -] - -[[package]] -name = "windows-sys" -version = "0.52.0" +name = "winapi" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" dependencies = [ - "windows-targets", + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", ] [[package]] -name = "windows-sys" -version = "0.59.0" +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets", -] +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] -name = "windows-targets" -version = "0.52.6" +name = "winapi-util" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows-sys", ] [[package]] -name = "windows_aarch64_gnullvm" -version = "0.52.6" +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] -name = "windows_aarch64_msvc" -version = "0.52.6" +name = "windows-link" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" [[package]] -name = "windows_i686_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" - -[[package]] -name = "windows_i686_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" - -[[package]] -name = "windows_i686_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" - -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" - -[[package]] -name = "wit-bindgen" -version = "0.46.0" +name = "windows-sys" +version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" dependencies = [ - "byteorder", "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.35" +version = "0.8.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/Cargo.toml b/Cargo.toml index d49c545..71b7325 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "blitztext" -version = "0.1.6" +version = "0.2.0" edition = "2021" authors = ["Praise Oketola "] description = "A library for fast keyword extraction and replacement in strings." @@ -14,16 +14,16 @@ name = "blitztext" crate-type = ["cdylib", "rlib"] [dependencies] -rand = "0.9.2" +rand = "0.10.3" regex = "1.12.2" rayon = "1.11.0" -pyo3 = { version = "0.22.6", features = ["extension-module"] } +pyo3 = { version = "0.29.3", features = ["extension-module"] } rustc-hash = "2.1.1" bincode = "1.3.3" serde = { version = "1.0.228", features = ["derive"] } [dev-dependencies] -criterion = "0.5.1" +criterion = "0.8.2" [[bench]] name = "benchmark" diff --git a/README.md b/README.md index 1465f6d..6648837 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,8 @@ replaced = processor.replace_keywords(text) print(f"Replaced text: {replaced}") ``` +`start` and `end` are UTF-8 byte offsets into the text, in Rust and Python alike. For non-ASCII text in Python, slice the encoded text: `text.encode()[m.start:m.end].decode()`. + ## Features ### 1. Parallel Processing @@ -111,10 +113,12 @@ texts = ["Text 1", "Text 2", "Text 3"] results = processor.parallel_extract_keywords_from_texts(texts) ``` +The parallel methods release the GIL while they run, so other Python threads keep going. Changing the same processor from another thread during such a call (for example `add_keyword`) raises `RuntimeError: Already borrowed`. + ### 2. Fuzzy Matching -Both Rust and Python implementations support fuzzy matching: +Both Rust and Python implementations support fuzzy matching. Similarity is `1 - edits / max(matched length, keyword length)`, counted in characters: ```rust // Rust diff --git a/blitztext.pyi b/blitztext.pyi index 1bf6f28..6317e69 100644 --- a/blitztext.pyi +++ b/blitztext.pyi @@ -7,8 +7,8 @@ class KeywordMatch: Attributes: keyword (str): The matched keyword or its clean name. similarity (float): The similarity score of the match (1.0 for exact matches). - start (int): The starting index of the match in the text. - end (int): The ending index of the match in the text. + start (int): UTF-8 byte offset where the match starts in the text. + end (int): UTF-8 byte offset where the match ends in the text. """ keyword: str similarity: float diff --git a/pyproject.toml b/pyproject.toml index d03874a..b553c3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "blitztext" -requires-python = ">=3.7" +requires-python = ">=3.8" classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", diff --git a/src/bin/manual_benchmark.rs b/src/bin/manual_benchmark.rs index d109e65..e1d03fc 100644 --- a/src/bin/manual_benchmark.rs +++ b/src/bin/manual_benchmark.rs @@ -3,7 +3,7 @@ use std::time::Instant; // use plotters::prelude::*; use blitztext::KeywordProcessor; use rand::seq::IndexedRandom; -use rand::Rng; +use rand::RngExt; // use regex::Regex; use regex::bytes::Regex; @@ -27,7 +27,7 @@ fn benchmark_keyword_extraction() -> Vec<(usize, f64, f64)> { for keywords_length in (0..=20000).step_by(1000) { let all_words_chosen = all_words - .choose_multiple(&mut rng, 5000) + .sample(&mut rng, 5000) .collect::>(); let story = all_words_chosen .iter() @@ -37,7 +37,7 @@ fn benchmark_keyword_extraction() -> Vec<(usize, f64, f64)> { // let story = all_words_chosen.iter().cloned().collect::>().join(" "); // let story = all_words_chosen.join(" "); let unique_keywords_sublist: Vec = all_words - .choose_multiple(&mut rng, keywords_length) + .sample(&mut rng, keywords_length) .cloned() .collect(); diff --git a/src/char_set.rs b/src/char_set.rs index ee4cdcb..66ca520 100644 --- a/src/char_set.rs +++ b/src/char_set.rs @@ -1,4 +1,5 @@ -use serde::{Deserialize, Serialize, Serializer, Deserializer}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use std::collections::BTreeSet; /// The size of the bitmap used to represent the character set. const BITMAP_SIZE: usize = (char::MAX as usize) / 64 + 1; @@ -144,3 +145,73 @@ impl CharSet { (self.bitmap[index / 64] & (1 << (index % 64))) != 0 } } + +/// A character set that stays small when it holds mostly ASCII and Latin-1 characters. +/// It serializes exactly like a `CharSet`. +#[derive(Clone, Serialize, Deserialize)] +#[serde(from = "CharSet", into = "CharSet")] +pub(crate) struct CompactCharSet { + latin1: [u64; 4], + others: BTreeSet, +} + +impl CompactCharSet { + pub fn new() -> Self { + CompactCharSet { + latin1: [0; 4], + others: BTreeSet::new(), + } + } + + pub fn insert(&mut self, c: char) { + let index = c as usize; + if index < 256 { + self.latin1[index / 64] |= 1 << (index % 64); + } else { + self.others.insert(c); + } + } + + pub fn extend(&mut self, iter: I) + where + I: IntoIterator, + { + for c in iter { + self.insert(c); + } + } + + pub fn contains(&self, c: char) -> bool { + let index = c as usize; + if index < 256 { + (self.latin1[index / 64] & (1 << (index % 64))) != 0 + } else { + self.others.contains(&c) + } + } +} + +impl From for CompactCharSet { + fn from(set: CharSet) -> Self { + let mut compact = CompactCharSet::new(); + compact.latin1.copy_from_slice(&set.bitmap[..4]); + for (word_index, &word) in set.bitmap.iter().enumerate().skip(4) { + let mut bits = word; + while bits != 0 { + let index = word_index * 64 + bits.trailing_zeros() as usize; + compact.others.extend(char::from_u32(index as u32)); + bits &= bits - 1; + } + } + compact + } +} + +impl From for CharSet { + fn from(compact: CompactCharSet) -> Self { + let mut set = CharSet::new(); + set.bitmap[..4].copy_from_slice(&compact.latin1); + set.extend(compact.others); + set + } +} diff --git a/src/compiled_trie.rs b/src/compiled_trie.rs new file mode 100644 index 0000000..5205a23 --- /dev/null +++ b/src/compiled_trie.rs @@ -0,0 +1,319 @@ +//! A read-only, cache-friendly copy of the keyword trie (`TrieNode`), used for matching. +//! +//! `KeywordProcessor` keeps the editable `TrieNode` tree as the source of truth (it is what +//! gets pickled) and builds a `CompiledTrie` from it lazily; any add or remove drops it. +//! +//! Layout: a double-array trie over byte classes. +//! - Each distinct byte used by a keyword gets a class 1..n; every other byte is class 0. +//! - State `s` has a slot `[base, check]`. Its child for byte `b` is `t = base[s] + class(b)`, +//! and that child exists only if `check[t]` equals `s` (low 31 bits). The top bit of +//! `check[t]` says whether `t` ends a keyword, so an exact walk reads one slot per byte. +//! - Free slots have `check == NONE`; the root is slot 0. Every `base + class` is kept in +//! bounds, so lookups never need a length check beyond the slice index. +//! +//! Building: states are placed breadth-first. For each state, `FreeSlots::find_base` walks a +//! linked list of free slots for a base where all its children fit. A slot that fails as a +//! candidate too often is dropped from the list, which bounds the search. +//! +//! Per state, `State` also keeps its child edges (for fuzzy search, which must visit every +//! child) and the distance in characters to its deepest keyword (the fuzzy pruning bound). + +use crate::fuzzy::TrieView; +use crate::trie::TrieNode; +use std::collections::VecDeque; + +const NONE: u32 = u32::MAX; +const KEYWORD_FLAG: u32 = 1 << 31; +const PARENT_MASK: u32 = !KEYWORD_FLAG; + +/// Read-only copy of the keyword trie, built for fast matching. +/// +/// States live in a double-array: the child of state `s` for byte `b` is +/// `t = base[s] + byte_classes[b]`, valid when `check[t]` names `s` as parent. +/// Fuzzy search also lists each state's children, so those are kept as edge lists. +pub(crate) struct CompiledTrie { + /// `[base, check]` per slot; `check` holds the parent slot and the keyword flag. + slots: Vec<[u32; 2]>, + /// Maps each byte to a class; bytes in no keyword share class 0, which never matches. + byte_classes: [u32; 256], + states: Vec, + edge_bytes: Vec, + edge_targets: Vec, + keywords: Vec>, +} + +#[derive(Clone, Copy)] +struct State { + first_edge: u32, + edge_count: u32, + keyword: u32, + /// Characters to the deepest keyword at or below this state, or NONE. + chars_to_deepest_keyword: u32, +} + +const EMPTY_STATE: State = State { + first_edge: 0, + edge_count: 0, + keyword: NONE, + chars_to_deepest_keyword: NONE, +}; + +/// Free double-array slots in a linked list, for finding bases while building. +#[derive(Default)] +struct FreeSlots { + next: Vec, + previous: Vec, + failures: Vec, + head: u32, + tail: u32, + linked: Vec, +} + +impl FreeSlots { + /// A slot that failed this many times as a base candidate is no longer tried. + const MAX_FAILURES: u8 = 16; + + /// Extends `slots` to `length` and links the new slots. + fn grow(&mut self, slots: &mut Vec<[u32; 2]>, length: usize) { + if self.linked.is_empty() { + self.head = NONE; + self.tail = NONE; + self.linked.push(false); + self.next.push(NONE); + self.previous.push(NONE); + self.failures.push(0); + } + slots.resize(length.max(slots.len()), [0, NONE]); + for (slot, entry) in slots.iter().enumerate().skip(self.linked.len()) { + let is_free = entry[1] == NONE; + self.linked.push(is_free); + self.next.push(NONE); + self.previous.push(if is_free { self.tail } else { NONE }); + self.failures.push(0); + if is_free { + match self.tail { + NONE => self.head = slot as u32, + tail => self.next[tail as usize] = slot as u32, + } + self.tail = slot as u32; + } + } + } + + fn unlink(&mut self, slot: usize) { + if !self.linked[slot] { + return; + } + self.linked[slot] = false; + let (previous, next) = (self.previous[slot], self.next[slot]); + match previous { + NONE => self.head = next, + previous => self.next[previous as usize] = next, + } + match next { + NONE => self.tail = previous, + next => self.previous[next as usize] = previous, + } + } + + /// Finds a base where every `base + class` slot is free. Grows `slots` so that + /// `base + class` stays in bounds for every class, matching or not. + fn find_base(&mut self, slots: &mut Vec<[u32; 2]>, classes: &[u32], class_count: u32) -> u32 { + let Some(&lowest_class) = classes.iter().min() else { + return 0; + }; + let mut candidate = self.head; + loop { + if candidate == NONE { + let old_tail = self.tail; + self.grow(slots, slots.len() + class_count as usize); + candidate = if old_tail == NONE { self.head } else { self.next[old_tail as usize] }; + continue; + } + let slot = candidate as usize; + if slot > lowest_class as usize { + let base = (slot - lowest_class as usize) as u32; + let needed = (base + class_count) as usize + 1; + if slots.len() < needed { + self.grow(slots, needed); + } + if classes.iter().all(|&class| slots[(base + class) as usize][1] == NONE) { + return base; + } + self.failures[slot] += 1; + if self.failures[slot] >= Self::MAX_FAILURES { + self.unlink(slot); + } + } + // Growing may have linked new slots after this one, so read the link now. + candidate = self.next[slot]; + } + } +} + +impl CompiledTrie { + pub fn new(root: &TrieNode) -> Self { + let mut byte_classes = [0u32; 256]; + let mut class_count = 1; + let mut pending = vec![root]; + while let Some(node) = pending.pop() { + for (&key, child) in &node.children { + let byte = key as u32 as u8; + if byte_classes[byte as usize] == 0 { + byte_classes[byte as usize] = class_count; + class_count += 1; + } + pending.push(child); + } + } + + let mut trie = CompiledTrie { + slots: vec![[0, NONE]; class_count as usize + 1], + byte_classes, + states: Vec::new(), + edge_bytes: Vec::new(), + edge_targets: Vec::new(), + keywords: Vec::new(), + }; + trie.slots[0][1] = PARENT_MASK; + let mut free_slots = FreeSlots::default(); + free_slots.grow(&mut trie.slots, class_count as usize + 1); + let mut order = Vec::new(); + let mut queue = VecDeque::from([(root, 0u32)]); + while let Some((node, slot)) = queue.pop_front() { + order.push(slot); + let classes: Vec = node + .children + .keys() + .map(|&key| byte_classes[key as u32 as u8 as usize]) + .collect(); + let base = free_slots.find_base(&mut trie.slots, &classes, class_count); + trie.slots[slot as usize][0] = base; + + if trie.states.len() <= slot as usize { + trie.states.resize(slot as usize + 1, EMPTY_STATE); + } + let state = &mut trie.states[slot as usize]; + state.first_edge = trie.edge_bytes.len() as u32; + state.edge_count = classes.len() as u32; + if let Some(keyword) = &node.keyword { + state.keyword = trie.keywords.len() as u32; + state.chars_to_deepest_keyword = 0; + trie.keywords.push(keyword.as_str().into()); + } + + for ((&key, child), class) in node.children.iter().zip(classes) { + let child_slot = base + class; + let keyword_flag = if child.keyword.is_some() { KEYWORD_FLAG } else { 0 }; + trie.slots[child_slot as usize][1] = slot | keyword_flag; + free_slots.unlink(child_slot as usize); + trie.edge_bytes.push(key as u32 as u8); + trie.edge_targets.push(child_slot); + queue.push_back((child, child_slot)); + } + } + trie.states.resize(trie.slots.len(), EMPTY_STATE); + + // Children come after their parent in breadth-first order, so walk it backwards. + for &slot in order.iter().rev() { + for edge in trie.edges(slot) { + let child = trie.states[trie.edge_targets[edge] as usize]; + if child.chars_to_deepest_keyword == NONE { + continue; + } + // UTF-8 continuation bytes do not start a new character. + let starts_char = trie.edge_bytes[edge] & 0xC0 != 0x80; + let chars = child.chars_to_deepest_keyword + starts_char as u32; + let parent = &mut trie.states[slot as usize]; + if parent.chars_to_deepest_keyword == NONE || parent.chars_to_deepest_keyword < chars { + parent.chars_to_deepest_keyword = chars; + } + } + } + trie + } + + /// Approximate heap and inline size in bytes. + pub fn memory_usage(&self) -> usize { + size_of::() + + self.slots.capacity() * size_of::<[u32; 2]>() + + self.states.capacity() * size_of::() + + self.edge_bytes.capacity() + + self.edge_targets.capacity() * size_of::() + + self.keywords.iter().map(|k| size_of::>() + k.len()).sum::() + } + + fn edges(&self, slot: u32) -> std::ops::Range { + let state = self.states[slot as usize]; + state.first_edge as usize..(state.first_edge + state.edge_count) as usize + } + + #[inline] + fn child(&self, slot: u32, byte: u8) -> Option { + let class = self.byte_classes[byte as usize]; + if class == 0 { + return None; + } + let child = self.slots[slot as usize][0] + class; + (self.slots[child as usize][1] & PARENT_MASK == slot).then_some(child) + } + + fn keyword(&self, slot: u32) -> Option<&str> { + let keyword = self.states[slot as usize].keyword; + (keyword != NONE).then(|| &*self.keywords[keyword as usize]) + } +} + +impl<'t> TrieView<'t> for &'t CompiledTrie { + type Node = u32; + + fn root(&self) -> Self::Node { + 0 + } + + fn node_id(&self, node: Self::Node) -> usize { + node as usize + } + + fn clean_name(&self, node: Self::Node) -> Option<&'t str> { + self.keyword(node) + } + + fn byte_child(&self, node: Self::Node, byte: u8) -> Option { + self.child(node, byte) + } + + fn for_each_byte_child(&self, node: Self::Node, mut visit: impl FnMut(u8, Self::Node)) { + for edge in self.edges(node) { + visit(self.edge_bytes[edge], self.edge_targets[edge]); + } + } + + fn chars_to_deepest_keyword(&self, node: Self::Node) -> Option { + let chars = self.states[node as usize].chars_to_deepest_keyword; + (chars != NONE).then_some(chars as usize) + } + + fn longest_prefix(&self, text: &[u8]) -> Option<(&'t str, usize)> { + let mut slot = 0; + let mut found = None; + for (i, &byte) in text.iter().enumerate() { + let class = self.byte_classes[byte as usize]; + if class == 0 { + break; + } + // Inlined `child()`: the same check word also carries the keyword flag. + let child = self.slots[slot as usize][0] + class; + let check = self.slots[child as usize][1]; + if check & PARENT_MASK != slot { + break; + } + slot = child; + if check & KEYWORD_FLAG != 0 { + found = Some((slot, i + 1)); + } + } + found.map(|(slot, length)| (self.keyword(slot).unwrap(), length)) + } + +} diff --git a/src/fuzzy.rs b/src/fuzzy.rs new file mode 100644 index 0000000..2c4563c --- /dev/null +++ b/src/fuzzy.rs @@ -0,0 +1,199 @@ +//! Fuzzy prefix search over a keyword trie, shared by the editable and the compiled trie. +//! +//! At each word start, `best_fuzzy_prefix` finds the keyword `k` and text span `w` (at least +//! one character) maximising `1 - edits(w, k) / max(|w|, |k|)` in characters, at or above the +//! threshold. Ties go to a span ending in a letter or digit, then the longer span, then the +//! smaller clean name. +//! +//! Search: depth-first over states (trie node, text characters consumed, edits so far). +//! Moves are exact (free), substitution, insertion (skip a text character) and deletion (skip a +//! keyword character), each one edit. Edit moves are taken only while +//! `1 - edits / (span + 1) >= threshold`; at the end of the text only deletions remain. +//! +//! The result does not depend on visit order, so two prunings leave it unchanged: +//! - Bound: a match below a state costs at least its edits and is normalised by at most +//! `max(span, keyword chars) + chars to the deepest keyword below`. +//! - Memo: a branching state seen again at the same (node, span) with no fewer edits can +//! only reproduce candidates found from the first visit. +//! +//! The exact child is pushed last, so it is explored first. + +use rustc_hash::FxHashMap; +use std::cmp::Reverse; + +/// Read access to a keyword trie, implemented by the editable and the compiled trie. +pub(crate) trait TrieView<'t>: Copy { + type Node: Copy; + + fn root(&self) -> Self::Node; + /// A key that identifies the node within this trie. + fn node_id(&self, node: Self::Node) -> usize; + /// The clean name stored at the node, if the node ends a keyword. + fn clean_name(&self, node: Self::Node) -> Option<&'t str>; + fn byte_child(&self, node: Self::Node, byte: u8) -> Option; + fn for_each_byte_child(&self, node: Self::Node, visit: impl FnMut(u8, Self::Node)); + /// An upper bound on the characters from this node to the deepest keyword below it, + /// or None when no keyword is at or below it. + fn chars_to_deepest_keyword(&self, node: Self::Node) -> Option; + /// Returns the longest keyword that `text` starts with, and its length in bytes. + fn longest_prefix(&self, text: &[u8]) -> Option<(&'t str, usize)>; + + fn char_child(&self, node: Self::Node, c: char) -> Option { + if c.is_ascii() { + return self.byte_child(node, c as u8); + } + let mut buffer = [0; 4]; + c.encode_utf8(&mut buffer) + .bytes() + .try_fold(node, |node, byte| self.byte_child(node, byte)) + } + + /// Visits each child one whole character below `node`, decoding multi-byte paths. + fn for_each_char_child(&self, node: Self::Node, mut visit: impl FnMut(char, Self::Node)) { + fn walk<'t, T: TrieView<'t>>( + trie: &T, + node: T::Node, + bytes: &mut ([u8; 4], usize), + visit: &mut impl FnMut(char, T::Node), + ) { + trie.for_each_byte_child(node, |byte, child| { + if bytes.1 == 0 && byte.is_ascii() { + visit(byte as char, child); + return; + } + bytes.0[bytes.1] = byte; + bytes.1 += 1; + match std::str::from_utf8(&bytes.0[..bytes.1]) { + Ok(text) => visit(text.chars().next().unwrap(), child), + Err(error) if error.error_len().is_none() => walk(trie, child, bytes, visit), + Err(_) => {} + } + bytes.1 -= 1; + }); + } + walk(self, node, &mut ([0; 4], 0), &mut visit); + } +} + +/// Buffers reused across the fuzzy searches of one text. +pub(crate) struct FuzzyScratch { + /// Search states: (node, text characters consumed, cost, keyword characters). + stack: Vec<(N, usize, usize, usize)>, + lowest_expanded_cost: FxHashMap<(usize, usize), usize>, + /// Characters of the text being searched, with the byte offset where each ends. + chars: Vec<(char, usize)>, +} + +impl Default for FuzzyScratch { + fn default() -> Self { + FuzzyScratch { + stack: Vec::new(), + lowest_expanded_cost: FxHashMap::default(), + chars: Vec::new(), + } + } +} + +struct Candidate<'t> { + clean_name: &'t str, + chars: usize, + bytes: usize, + similarity: f32, + ends_alphanumeric: bool, +} + +impl Candidate<'_> { + /// Higher similarity wins, then a span ending in a letter or digit, then the longer span, + /// then the smaller clean name. + fn beats(&self, other: &Candidate) -> bool { + let rank = |c: &Candidate| (c.similarity, c.ends_alphanumeric, c.chars); + (rank(self), Reverse(self.clean_name)) > (rank(other), Reverse(other.clean_name)) + } +} + +/// Finds the keyword most similar to a prefix of `text`: `1 - edits / max(span, keyword length)` +/// in characters. An edit is allowed only while `1 - edits / (span + 1)` stays at `threshold`. +/// Returns (clean name, matched length in bytes, similarity). +pub(crate) fn best_fuzzy_prefix<'t, T: TrieView<'t>>( + trie: T, + text: &str, + threshold: f32, + scratch: &mut FuzzyScratch, +) -> (Option<&'t str>, usize, f32) { + let FuzzyScratch { stack, lowest_expanded_cost, chars } = scratch; + lowest_expanded_cost.clear(); + chars.clear(); + let mut unread = text.char_indices(); + let mut char_at = |index: usize, chars: &mut Vec<(char, usize)>| { + while chars.len() <= index { + let (offset, c) = unread.next()?; + chars.push((c, offset + c.len_utf8())); + } + Some(chars[index].0) + }; + + let mut best: Option> = None; + // Every match below a state costs at least `cost` and is normalised by at most + // `max(span, keyword length) + chars to the deepest keyword`. + let can_beat = |node: T::Node, span: usize, cost: usize, keyword_chars: usize, best: &Option| { + let Some(remaining) = trie.chars_to_deepest_keyword(node) else { + return false; + }; + let longest = span.max(keyword_chars).saturating_add(remaining); + let best_possible = if cost == 0 { 1.0 } else { 1.0 - (cost as f32 / longest as f32) }; + best_possible >= threshold && best.as_ref().is_none_or(|best| best_possible >= best.similarity) + }; + + stack.clear(); + stack.push((trie.root(), 0, 0, 0)); + while let Some((node, span, cost, keyword_chars)) = stack.pop() { + if !can_beat(node, span, cost, keyword_chars, &best) { + continue; + } + if let (Some(clean_name), true) = (trie.clean_name(node), span > 0) { + let similarity = 1.0 - (cost as f32 / span.max(keyword_chars) as f32); + let candidate = Candidate { + clean_name, + chars: span, + bytes: chars[span - 1].1, + similarity, + ends_alphanumeric: chars[span - 1].0.is_alphanumeric(), + }; + if similarity >= threshold && best.as_ref().is_none_or(|best| candidate.beats(best)) { + best = Some(candidate); + } + } + + // None once the text is used up; only deletions remain possible then. + let current_char = char_at(span, chars); + let can_branch = 1.0 - (cost as f32 / (span + 1) as f32) >= threshold; + if can_branch { + // A state reached again at no lower cost yields only matches that are no better. + let lowest = lowest_expanded_cost.entry((trie.node_id(node), span)).or_insert(usize::MAX); + if *lowest <= cost { + continue; + } + *lowest = cost; + // Insertion: skip the current text character + if current_char.is_some() { + stack.push((node, span + 1, cost + 1, keyword_chars)); + } + trie.for_each_char_child(node, |c, child| { + if current_char.is_some_and(|current| c != current) { + // Substitution + stack.push((child, span + 1, cost + 1, keyword_chars + 1)); + } + // Deletion: skip the keyword character + stack.push((child, span, cost + 1, keyword_chars + 1)); + }); + } + // Pushed last so the exact path is searched first and finds strong matches early. + if let Some(child) = current_char.and_then(|current| trie.char_child(node, current)) { + stack.push((child, span + 1, cost, keyword_chars + 1)); + } + } + match best { + Some(best) => (Some(best.clean_name), best.bytes, best.similarity), + None => (None, 0, 0.0), + } +} diff --git a/src/keyword_processor.rs b/src/keyword_processor.rs index 63f2749..110580a 100644 --- a/src/keyword_processor.rs +++ b/src/keyword_processor.rs @@ -1,25 +1,37 @@ -use crate::char_set::CharSet; +use crate::char_set::CompactCharSet; +use crate::compiled_trie::CompiledTrie; +use crate::fuzzy::{self, FuzzyScratch, TrieView}; use crate::trie::TrieNode; use crate::utils::{is_whitespace, round}; use rayon::prelude::*; -use std::cmp::{max, min}; use std::collections::HashMap; use std::mem::size_of; +use std::borrow::Cow; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::OnceLock; -use std::sync::Arc; use serde::{Deserialize, Serialize}; /// A struct for handling keyword matching, fuzzy matching, and replacement in text. #[derive(Serialize, Deserialize)] pub struct KeywordProcessor { /// Set of characters that are considered part of a word. - non_word_boundaries: CharSet, + non_word_boundaries: CompactCharSet, /// The root of the keyword trie. root: TrieNode, /// Whether the keyword matching should be case-sensitive. case_sensitive: bool, /// Whether to allow overlapping matches. allow_overlaps: bool, + /// Faster read-only copy of `root`, dropped when keywords change. + #[serde(skip)] + compiled: OnceLock, + /// Number of nodes in `root`; counted on first use after deserializing. + #[serde(skip)] + node_count: OnceLock, + /// Matching work done with `root` since keywords last changed, in text bytes. + #[serde(skip)] + work_since_change: AtomicUsize, } impl KeywordProcessor { @@ -61,7 +73,7 @@ impl KeywordProcessor { /// let case_sensitive_processor = KeywordProcessor::with_options(true, false); /// ``` pub fn with_options(case_sensitive: bool, allow_overlaps: bool) -> Self { - let mut non_word_boundaries = CharSet::new(); + let mut non_word_boundaries = CompactCharSet::new(); non_word_boundaries.extend('a'..='z'); non_word_boundaries.extend('A'..='Z'); non_word_boundaries.extend('0'..='9'); @@ -71,6 +83,9 @@ impl KeywordProcessor { root: TrieNode::new(), case_sensitive, allow_overlaps, + compiled: OnceLock::new(), + node_count: OnceLock::new(), + work_since_change: AtomicUsize::new(0), } } @@ -116,15 +131,25 @@ impl KeywordProcessor { let clean_name = clean_name.unwrap_or(&keyword).to_string(); + let mut added_nodes = 0; let mut current_node = &mut self.root; // we use bytes as char instead of char because extract_keywords uses bytes for byte in keyword.bytes() { - current_node = current_node - .children - .entry(byte as char) - .or_insert_with(|| Box::new(TrieNode::new())); + current_node = current_node.children.entry(byte as char).or_insert_with(|| { + added_nodes += 1; + Box::new(TrieNode::new()) + }); } current_node.keyword = Some(clean_name); + self.keywords_changed(added_nodes as isize); + } + + fn keywords_changed(&mut self, node_count_change: isize) { + self.compiled.take(); + *self.work_since_change.get_mut() = 0; + if let Some(node_count) = self.node_count.get_mut() { + *node_count = node_count.saturating_add_signed(node_count_change); + } } /// Removes a keyword from the processor. @@ -157,7 +182,12 @@ impl KeywordProcessor { keyword.to_lowercase() }; - Self::remove_keyword_helper(&keyword, &mut self.root) + let mut removed_nodes = 0; + let removed = Self::remove_keyword_helper(keyword.as_bytes(), &mut self.root, &mut removed_nodes); + if removed { + self.keywords_changed(-(removed_nodes as isize)); + } + removed } /// Helper function for removing a keyword from the trie. @@ -194,25 +224,21 @@ impl KeywordProcessor { /// * Time complexity: O(m), where m is the length of the keyword. /// * Space complexity: O(m) due to the recursion stack. /// - fn remove_keyword_helper(keyword: &str, node: &mut TrieNode) -> bool { - if keyword.is_empty() { - if node.keyword.is_some() { - node.keyword = None; - true - } else { - false - } - } else { - let (first_char, rest) = keyword.split_at(1); - if let Some(child) = node.children.get_mut(&first_char.chars().next().unwrap()) { - let removed = Self::remove_keyword_helper(rest, child); + fn remove_keyword_helper(keyword: &[u8], node: &mut TrieNode, removed_nodes: &mut usize) -> bool { + if let Some((&first_byte, rest)) = keyword.split_first() { + let key = first_byte as char; + if let Some(child) = node.children.get_mut(&key) { + let removed = Self::remove_keyword_helper(rest, child, removed_nodes); if removed && child.children.is_empty() && child.keyword.is_none() { - node.children.remove(&first_char.chars().next().unwrap()); + node.children.remove(&key); + *removed_nodes += 1; } removed } else { false } + } else { + node.keyword.take().is_some() } } @@ -300,7 +326,7 @@ impl KeywordProcessor { /// ``` pub fn get_all_keywords(&self) -> HashMap { let mut keywords = HashMap::new(); - self.collect_keywords(&self.root, String::new(), &mut keywords); + self.collect_keywords(&self.root, &mut Vec::new(), &mut keywords); keywords } @@ -323,16 +349,17 @@ impl KeywordProcessor { fn collect_keywords( &self, node: &TrieNode, - prefix: String, + prefix: &mut Vec, keywords: &mut HashMap, ) { if let Some(clean_name) = &node.keyword { - keywords.insert(prefix.clone(), clean_name.clone()); + let keyword = String::from_utf8_lossy(prefix).into_owned(); + keywords.insert(keyword, clean_name.clone()); } - for (ch, child) in &node.children { - let mut new_prefix = prefix.clone(); - new_prefix.push(*ch); - self.collect_keywords(child, new_prefix, keywords); + for (&byte, child) in &node.children { + prefix.push(byte as u8); + self.collect_keywords(child, prefix, keywords); + prefix.pop(); } } @@ -354,7 +381,7 @@ impl KeywordProcessor { /// processor.set_non_word_boundaries(&['a', 'b', 'c', '1', '2', '3']); /// ``` pub fn set_non_word_boundaries(&mut self, boundaries: &[char]) { - self.non_word_boundaries = CharSet::new(); + self.non_word_boundaries = CompactCharSet::new(); self.non_word_boundaries.extend(boundaries.iter().copied()); } @@ -438,19 +465,10 @@ impl KeywordProcessor { /// assert_eq!(processor.get_keyword("python"), None); /// ``` pub fn get_keyword(&self, keyword: &str) -> Option { - let keyword = if self.case_sensitive { - keyword.to_string() - } else { - keyword.to_lowercase() - }; - + let keyword = self.normalize_case(keyword); let mut current_node = &self.root; - for ch in keyword.chars() { - if let Some(next_node) = current_node.children.get(&ch) { - current_node = next_node; - } else { - return None; - } + for byte in keyword.bytes() { + current_node = current_node.children.get(&(byte as char))?; } current_node.keyword.clone() } @@ -497,87 +515,12 @@ impl KeywordProcessor { "Threshold must be between 0 and 1" ); - // Initialize vector to store found matches + let original = text; + let text = self.normalize_case(original); let mut matches = Vec::new(); - - let text = if self.case_sensitive { - text.to_string() - } else { - text.to_lowercase() - }; - - // working with bytes is faster than strings - let text_bytes = text.as_bytes(); - - // Iterate through each character in the text - let mut i = 0; - while i < text_bytes.len() { - // Check if we're at a word boundary - // This is true if we're at the start of the text or the previous character is not a word character - if i == 0 || !self.is_non_word_boundary(text_bytes[i - 1] as char) { - // Choose between fuzzy and exact matching based on threshold - // let (found, length, similarity) = self.find_keyword(&text_bytes[i..]); - let (found, length, similarity) = if threshold < 1.0 { - self.find_keyword_fuzzy(&text_bytes[i..], threshold) - } else { - self.find_keyword(&text_bytes[i..]) - }; - - // If a keyword is found - if let Some(keyword) = found { - if threshold == 1.0 { - // For exact matching, keep the word boundary check - // For ensuring that we only match complete words, not partial words or substrings within larger words. - // This is true if we're at the end of the text or the next character is not a word character - if i + length == text_bytes.len() - || !self.is_non_word_boundary(text_bytes[i + length] as char) - { - matches.push(KeywordMatch { - keyword, - similarity: round(similarity, 2), - start: if is_whitespace(text_bytes[i]) { - i + 1 - } else { - i - }, - end: i + length, - }); - // Move the index past this keyword - i += length; - continue; - } - } else { - // For fuzzy matching, add the match with lenient boundary check - // on of next 3 characters is a non-word boundary - let end = if is_whitespace(text_bytes[i + length - 1]) { - i + length - 1 - } else { - i + length - }; - let is_boundary_in_next_x_chars = (end..min(end + 3, text.len())) - .any(|j| !self.is_non_word_boundary(text_bytes[j] as char)); - if is_boundary_in_next_x_chars || end == text_bytes.len() { - matches.push(KeywordMatch { - keyword, - similarity: round(similarity, 2), - start: if is_whitespace(text_bytes[i]) { - i + 1 - } else { - i - }, - end, - }); - // Move the index past this keyword - i += length; - continue; - } - } - } - } - // If no keyword found or not at a word boundary, move to the next character - i += 1; - } - + let trie = self.compiled_trie(text.len(), threshold); + self.scan_with(trie, &text, 0..text.len(), threshold, &mut matches); + self.map_to_original_offsets(original, &mut matches); matches } @@ -624,85 +567,23 @@ impl KeywordProcessor { "Threshold must be between 0 and 1" ); - let text = if self.case_sensitive { - text.to_string() - } else { - text.to_lowercase() - }; - - let text_bytes = text.as_bytes(); - let processor = Arc::new(self); - - let chunk_size = (text_bytes.len() / rayon::current_num_threads()).max(1000); + let original = text; + let text = self.normalize_case(original); + let trie = self.compiled_trie(text.len(), threshold); + let chunk_size = (text.len() / rayon::current_num_threads()).max(1000); - (0..text_bytes.len()) + let mut matches: Vec = (0..text.len()) .into_par_iter() .step_by(chunk_size) .flat_map(|start| { - let end = (start + chunk_size).min(text_bytes.len()); - let processor = Arc::clone(&processor); - + let end = (start + chunk_size).min(text.len()); let mut local_matches = Vec::new(); - let mut i = start; - while i < end { - if i == 0 || !processor.is_non_word_boundary(text_bytes[i - 1] as char) { - let (found, length, similarity) = if threshold < 1.0 { - processor.find_keyword_fuzzy(&text_bytes[i..], threshold) - } else { - processor.find_keyword(&text_bytes[i..]) - }; - - if let Some(keyword) = found { - if threshold == 1.0 { - if i + length == text_bytes.len() - || !processor - .is_non_word_boundary(text_bytes[i + length] as char) - { - local_matches.push(KeywordMatch { - keyword, - similarity: round(similarity, 2), - start: if is_whitespace(text_bytes[i]) { - i + 1 - } else { - i - }, - end: i + length, - }); - i += length; - continue; - } - } else { - let end = if is_whitespace(text_bytes[i + length - 1]) { - i + length - 1 - } else { - i + length - }; - let is_boundary_in_next_x_chars = - (end..min(end + 3, text_bytes.len())).any(|j| { - !processor.is_non_word_boundary(text_bytes[j] as char) - }); - if is_boundary_in_next_x_chars || end == text_bytes.len() { - local_matches.push(KeywordMatch { - keyword, - similarity: round(similarity, 2), - start: if is_whitespace(text_bytes[i]) { - i + 1 - } else { - i - }, - end, - }); - i += length; - continue; - } - } - } - } - i += 1; - } + self.scan_with(trie, &text, start..end, threshold, &mut local_matches); local_matches }) - .collect() + .collect(); + self.map_to_original_offsets(original, &mut matches); + matches } /// Extracts keywords from multiple texts in parallel. @@ -753,135 +634,164 @@ impl KeywordProcessor { .collect() } - /// Helper method to find a keyword starting from a given position in the text. - /// - /// This method traverses the trie to find an exact match for a keyword in the text. - /// - /// # Arguments - /// - /// * `text_bytes` - The text to search in, as a byte slice. - /// - /// # Returns - /// - /// A tuple containing an Option with the found keyword (if any), the length of the match, and the similarity (1.0 for exact matches). - /// - /// # Complexity - /// - /// * Time complexity: O(m), where m is the length of the matched keyword. - fn find_keyword(&self, text_bytes: &[u8]) -> (Option<&str>, usize, f32) { - // covert text_bytes to string - let mut current_node = &self.root; - let mut last_keyword_node = None; // Last node representing a keyword - let mut last_keyword_pos = 0; // Position of the last keyword found - - // for each character in the text starting from the current position - // check if there is a matching child node in the trie - for (i, &byte) in text_bytes.iter().enumerate() { - // Move to the next node in the trie - // println!("byte: {:?}, char: {:?}, is_byte_valid_char: {:?}", byte, byte as char, is_valid_char(&[byte])); - if let Some(next_node) = ¤t_node.children.get(&(byte as char)) { - // if let Some(next_node) = ¤t_node.children[byte as usize] { - current_node = next_node; - if current_node.keyword.is_some() { - // If this node represents a keyword, remember it - last_keyword_node = Some(current_node); - last_keyword_pos = i; - } - } else { - break; - } - } - // println!("last_keyword_node: {:?}, last_keyword_pos: {:?}", last_keyword_node, last_keyword_pos); - if let Some(node) = last_keyword_node { - (node.keyword.as_deref(), last_keyword_pos + 1, 1.0) - } else { - (None, 0, 0.0) + /// Scans with the compiled trie when there is one, else with the editable trie. + fn scan_with<'a>( + &'a self, + trie: Option<&'a CompiledTrie>, + text: &str, + positions: std::ops::Range, + threshold: f32, + matches: &mut Vec>, + ) { + match trie { + Some(trie) => self.scan(trie, text, positions, threshold, matches), + None => self.scan(&self.root, text, positions, threshold, matches), } } - /// Helper method to find a keyword starting from a given position in the text with fuzzy matching. - /// - /// This method uses a depth-first search through the trie to find approximate matches, - /// allowing for insertions, deletions, and substitutions up to the specified threshold. - /// - /// # Arguments - /// - /// * `text_bytes` - The text to search in, as a byte slice. - /// * `threshold` - The minimum similarity required for a match (0.0 to 1.0). - /// - /// # Returns - /// - /// A tuple containing an Option with the found keyword (if any), the length of the match, and the similarity. - /// - /// # Complexity - /// - /// * Time complexity: O(n * m), where n is the length of the text and m is the average length of keywords. - /// In practice, the threshold and early termination optimizations make it much faster for most inputs. - fn find_keyword_fuzzy(&self, text_bytes: &[u8], threshold: f32) -> (Option<&str>, usize, f32) { - // Initialize variables to keep track of the best match found - // let mut best_match: Option = None; - let mut best_match: Option<&str> = None; - let mut best_match_length = 0; - let mut best_similarity = 0.0; - - // Stack for depth-first search through the trie - // Each element is (node, depth in text, cumulative cost) - let mut stack = vec![(&self.root, 0, 0)]; - - // Depth-first search through the trie - while let Some((node, depth, cost)) = stack.pop() { - // Check if this node represents a keyword - if let Some(keyword) = &node.keyword { - let keyword_len = keyword.len(); - // Use the maximum of current depth and keyword length for normalization - let max_length = max(depth, keyword_len); - // Calculate similarity as 1 - (cost / max_length) - let similarity = 1.0 - (cost as f32 / max_length as f32); - // println!("similarity: {:?}, depth: {:?}, cost: {:?}, s>=t: {:?}, best_similarity: {:?}, best_match: {:?}", similarity, depth, cost, similarity >= threshold, best_similarity, best_match); - - // Update the best match if the similarity is above the threshold - if similarity >= threshold && similarity > best_similarity { - best_match = Some(keyword); - // best_match = Some(keyword.clone()); - best_match_length = depth; - best_similarity = similarity; - } + /// Appends the matches that start inside `positions` to `matches`; a match may run past its end. + fn scan<'a, T: TrieView<'a>>( + &self, + trie: T, + text: &str, + positions: std::ops::Range, + threshold: f32, + matches: &mut Vec>, + ) { + let text_bytes = text.as_bytes(); + let is_ascii_word_byte: [bool; 128] = + std::array::from_fn(|byte| self.is_non_word_boundary(byte as u8 as char)); + // Whether the character at byte `i` is a word character, and its length in bytes + let word_char_at = |i: usize| match text_bytes[i] { + byte if byte < 128 => (is_ascii_word_byte[byte as usize], 1), + _ => { + let c = text[i..].chars().next().unwrap(); + (self.is_non_word_boundary(c), c.len_utf8()) } + }; + let is_word_char_before = |i: usize| match text_bytes[i - 1] { + byte if byte < 128 => is_ascii_word_byte[byte as usize], + _ => self.is_non_word_boundary(text[..i].chars().next_back().unwrap()), + }; - // Stop exploring this path if we've reached the end of the text, or it's word boundary - if depth >= text_bytes.len() { - // if depth >= text_bytes.len() || (depth >5 && !self.is_non_word_boundary(text_bytes[depth-5] as char)) { + let mut fuzzy_scratch = FuzzyScratch::default(); + let mut i = positions.start; + 'scan: while i < positions.end { + if !text.is_char_boundary(i) { + i += 1; continue; } - - // Get the current byte(char) in the text - let current_char = text_bytes[depth] as char; - - // Explore exact matches (no additional cost) - if let Some(child) = &node.children.get(¤t_char) { - stack.push((child, depth + 1, cost)); + // A keyword can only start at the text start or after a non-word character + if i > 0 && is_word_char_before(i) { + loop { + let Some(&byte) = text_bytes[..positions.end].get(i) else { + break 'scan; + }; + // Runs once per text byte, so ASCII stays inline + let (is_word, char_len) = match byte { + byte if byte < 128 => (is_ascii_word_byte[byte as usize], 1), + _ => word_char_at(i), + }; + i += char_len; + if !is_word { + continue 'scan; + } + } } - - // Explore fuzzy matches - // Calculate current similarity to determine if we should continue exploring - let current_similarity = 1.0 - (cost as f32 / (depth + 1) as f32); - if current_similarity >= threshold { - // Insertion: skip the current character in the text (cost + 1) - stack.push((node, depth + 1, cost + 1)); - - // Deletion and Substitution - for (ch, child) in &node.children { - if *ch != current_char { - // Substitution: try all children with increased cost (cost + 1) - stack.push((child, depth + 1, cost + 1)); + // Choose between fuzzy and exact matching based on threshold + let (found, length, similarity) = if threshold < 1.0 { + fuzzy::best_fuzzy_prefix(trie, &text[i..], threshold, &mut fuzzy_scratch) + } else { + match trie.longest_prefix(&text_bytes[i..]) { + Some((keyword, length)) => (Some(keyword), length, 1.0), + None => (None, 0, 0.0), + } + }; + + if let Some(keyword) = found { + let start = if is_whitespace(text_bytes[i]) { i + 1 } else { i }; + if threshold == 1.0 { + // For ensuring that we only match complete words, not partial words or substrings within larger words. + // This is true if we're at the end of the text or the next character is not a word character + if i + length == text_bytes.len() || !word_char_at(i + length).0 { + matches.push(KeywordMatch { + keyword, + similarity: round(similarity, 2), + start, + end: i + length, + }); + i += length; + continue; + } + } else { + // For fuzzy matching, accept a word boundary within the next 3 characters + let end = if is_whitespace(text_bytes[i + length - 1]) { + i + length - 1 + } else { + i + length + }; + let is_boundary_in_next_x_chars = text[end..] + .chars() + .take(3) + .any(|c| !self.is_non_word_boundary(c)); + // A lone whitespace match trims to an inverted span, which is no match. + if start <= end && (is_boundary_in_next_x_chars || end == text_bytes.len()) { + matches.push(KeywordMatch { + keyword, + similarity: round(similarity, 2), + start, + end, + }); + i += length; + continue; } - // Deletion: skip the current node in the trie (cost + 1) - stack.push((child, depth, cost + 1)); } } + // If no keyword found or not at a word boundary, move to the next character + i += 1; + } + } + + /// Matches are found in the lowercased text; when lowercasing changed its length, + /// moves each match onto the original characters it came from. + fn map_to_original_offsets(&self, original: &str, matches: &mut [KeywordMatch]) { + if self.case_sensitive || matches.is_empty() { + return; + } + let Some(char_starts) = lowercase_char_starts(original) else { + return; + }; + for keyword_match in matches { + let start = char_starts[keyword_match.start]; + keyword_match.end = if keyword_match.end == keyword_match.start { + start + } else { + let last_char = char_starts[keyword_match.end - 1]; + last_char + original[last_char..].chars().next().map_or(0, char::len_utf8) + }; + keyword_match.start = start; + } + } + + fn normalize_case<'t>(&self, text: &'t str) -> Cow<'t, str> { + if self.case_sensitive { + Cow::Borrowed(text) + } else { + Cow::Owned(text.to_lowercase()) } - // println!("best_match: {:?}, best_match_length: {:?}, best_similarity: {:?}", best_match, best_match_length, best_similarity); - (best_match, best_match_length, best_similarity) + } + + /// Returns the compiled trie once bytes scanned since the last edit (fuzzy weighted 64x) + /// reach 16x the node count; until then matching uses the editable trie. + fn compiled_trie(&self, text_len: usize, threshold: f32) -> Option<&CompiledTrie> { + if let Some(trie) = self.compiled.get() { + return Some(trie); + } + // Fuzzy matching costs far more per byte than exact matching. + let work = if threshold < 1.0 { text_len * 64 } else { text_len }; + let total_work = self.work_since_change.fetch_add(work, Ordering::Relaxed) + work; + let node_count = *self.node_count.get_or_init(|| self.root.count_nodes()); + (total_work >= node_count * 16).then(|| self.compiled.get_or_init(|| CompiledTrie::new(&self.root))) } /// Replaces keywords in the given text with their clean names. @@ -915,19 +825,16 @@ impl KeywordProcessor { /// assert_eq!(result, "I love Rust Programming Language programming"); /// ``` pub fn replace_keywords(&self, text: &str, threshold: Option) -> String { - let matches = self.extract_keywords(text, threshold); - let mut result = text.to_string(); - - // Sort matches by start position in descending order - let mut sorted_matches = matches; - sorted_matches.sort_by(|a, b| b.start.cmp(&a.start)); - - for keyword_match in sorted_matches { - let start = keyword_match.start; - let end = keyword_match.end; - result.replace_range(start..end, &keyword_match.keyword); + let mut result = String::with_capacity(text.len()); + let mut copied_up_to = 0; + for keyword_match in self.extract_keywords(text, threshold) { + // Two matches inside one character that lowercases to several can overlap. + let start = keyword_match.start.max(copied_up_to); + result.push_str(&text[copied_up_to..start]); + result.push_str(keyword_match.keyword); + copied_up_to = keyword_match.end.max(start); } - + result.push_str(&text[copied_up_to..]); result } @@ -990,16 +897,22 @@ impl KeywordProcessor { if self.allow_overlaps { return false; } - // Check if the character is in the set of non-word boundaries - self.non_word_boundaries.contains(ch) - // self.non_word_boundaries.contains(&ch) + if self.non_word_boundaries.contains(ch) { + return true; + } + // Text is lowercased before matching, so an uppercase boundary covers its lowercase form. + let mut uppercase = ch.to_uppercase(); + match (self.case_sensitive, uppercase.next(), uppercase.next()) { + (false, Some(upper), None) if upper != ch => self.non_word_boundaries.contains(upper), + _ => false, + } } /// Calculates the approximate memory usage of the trie in bytes. /// /// This method provides an estimate of the memory used by the trie structure, - /// including all nodes and the strings stored in them. It does not include - /// the memory used by the KeywordProcessor struct itself. + /// including all nodes and the strings stored in them, plus the compiled + /// matching copy once it has been built. /// /// # Returns /// @@ -1007,8 +920,9 @@ impl KeywordProcessor { pub fn calculate_memory_usage(&self) -> usize { let trie_memory = self.calculate_node_memory(&self.root); + let compiled_memory = self.compiled.get().map_or(0, CompiledTrie::memory_usage); // Add size of KeywordProcessor struct itself - trie_memory + size_of::() + trie_memory + compiled_memory + size_of::() } fn calculate_node_memory(&self, node: &TrieNode) -> usize { @@ -1032,6 +946,21 @@ impl KeywordProcessor { } } +/// For text whose lowercase form has a different byte length, maps every byte offset +/// of `text.to_lowercase()` to the start of the original character it came from. +fn lowercase_char_starts(text: &str) -> Option> { + let lowercase_len = |c: char| c.to_lowercase().map(char::len_utf8).sum::(); + if text.is_ascii() || text.chars().all(|c| lowercase_len(c) == c.len_utf8()) { + return None; + } + let mut starts = Vec::with_capacity(text.len() + 1); + for (offset, c) in text.char_indices() { + starts.extend(std::iter::repeat_n(offset, lowercase_len(c))); + } + starts.push(text.len()); + Some(starts) +} + /// A struct representing a keyword match in a text. #[derive(Debug, PartialEq)] pub struct KeywordMatch<'a> { @@ -1291,7 +1220,7 @@ mod tests { }, KeywordMatch { keyword: "coding", - similarity: 0.78, + similarity: 0.82, start: 20, end: 29, }, @@ -1516,4 +1445,166 @@ mod tests { // println!("{:?}", keywords); assert_eq!(keywords.len(), 0); } + + #[test] + fn compiled_trie_matches_like_the_editable_trie() { + use rand::{rngs::StdRng, RngExt, SeedableRng}; + let mut rng = StdRng::seed_from_u64(7); + let pieces = ["a", "b", "ab", "ba", "c", "é", "日", "😀", " ", " ", "-", "x"]; + let random_string = |rng: &mut StdRng, max_pieces: usize| -> String { + (0..rng.random_range(0..=max_pieces)) + .map(|_| pieces[rng.random_range(0..pieces.len())]) + .collect() + }; + for case in 0..300 { + let mut kp = KeywordProcessor::with_options(case % 3 == 0, case % 5 == 0); + for _ in 0..rng.random_range(1..40) { + let keyword = random_string(&mut rng, 4); + let clean_name = random_string(&mut rng, 3); + kp.add_keyword(&keyword, rng.random_bool(0.5).then_some(clean_name.as_str())); + } + let compiled = CompiledTrie::new(&kp.root); + for _ in 0..20 { + let text = random_string(&mut rng, 30); + for threshold in [1.0, 0.95, 0.8, 0.7, 0.6, 0.5] { + let (mut expected, mut actual) = (Vec::new(), Vec::new()); + kp.scan(&kp.root, &text, 0..text.len(), threshold, &mut expected); + kp.scan(&compiled, &text, 0..text.len(), threshold, &mut actual); + assert_eq!(actual, expected, "text {text:?} threshold {threshold}"); + } + } + } + } + + #[test] + fn non_ascii_non_word_boundary_joins_words() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("caf", None); + assert_eq!(kp.extract_keywords("caf café", None).len(), 2); + kp.add_non_word_boundary('é'); + let matches = kp.extract_keywords("caf café", None); + assert_eq!((matches.len(), matches[0].start, matches[0].end), (1, 0, 3)); + } + + #[test] + fn uppercase_boundary_applies_in_case_insensitive_mode() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("caf", None); + kp.add_non_word_boundary('É'); + assert_eq!(kp.extract_keywords("CAFÉ", None), vec![]); + assert_eq!(kp.extract_keywords("CAF.", None).len(), 1); + } + + #[test] + fn offsets_point_into_original_text_when_lowercase_is_longer() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("istanbul", Some("IST")); + let text = "İstanbul istanbul"; + let matches = kp.extract_keywords(text, None); + assert_eq!(&text[matches[0].start..matches[0].end], "istanbul"); + } + + #[test] + fn fuzzy_tie_goes_to_smallest_clean_name() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("cat", None); + kp.add_keyword("car", None); + let matches = kp.extract_keywords("cab", Some(0.6)); + assert_eq!((matches[0].keyword, matches[0].similarity), ("car", 0.67)); + } + + #[test] + fn fuzzy_matches_a_missing_last_letter_at_text_end() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("python", None); + let matches = kp.extract_keywords("I use pytho", Some(0.8)); + assert_eq!((matches[0].keyword, matches[0].start, matches[0].end), ("python", 6, 11)); + } + + #[test] + fn fuzzy_match_leaves_trailing_punctuation_out() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("python", None); + assert_eq!(kp.replace_keywords("I like pytho.", Some(0.8)), "I like python."); + } + + #[test] + fn fuzzy_whitespace_only_match_is_dropped() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("a ", Some("X")); + for text in ["Klong ", "klong "] { + assert!(kp.extract_keywords(text, Some(0.5)).iter().all(|m| m.start <= m.end)); + assert_eq!(kp.replace_keywords(text, Some(0.5)), text); + } + } + + #[test] + fn fuzzy_counts_a_wrong_accent_as_one_edit() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("cafe", None); + let matches = kp.extract_keywords("café", Some(0.75)); + assert_eq!((matches[0].similarity, matches[0].end), (0.75, 5)); + } + + fn processor_with_long_clean_name() -> KeywordProcessor { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("a", Some("")); + kp + } + + #[test] + fn fuzzy_extract_skips_zero_length_match_at_text_start() { + let kp = processor_with_long_clean_name(); + assert_eq!(kp.extract_keywords("b", Some(0.8)), vec![]); + assert_eq!(kp.replace_keywords("b", Some(0.8)), "b"); + } + + #[test] + fn fuzzy_extract_terminates_on_zero_length_match_after_text_start() { + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let kp = processor_with_long_clean_name(); + let matches: Vec<(usize, usize)> = kp + .extract_keywords("a b", Some(0.8)) + .iter() + .map(|m| (m.start, m.end)) + .collect(); + sender.send(matches).unwrap(); + }); + let matches = receiver + .recv_timeout(std::time::Duration::from_secs(10)) + .expect("fuzzy extraction did not finish"); + assert_eq!(matches, vec![(0, 1)]); + } + + #[test] + fn fuzzy_matches_start_and_end_on_character_boundaries() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("a", None); + kp.add_keyword("ab", None); + let matches = kp.extract_keywords("a😀b", Some(0.5)); + assert_eq!((matches[1].start, matches[1].end), (5, 6)); + assert_eq!(kp.replace_keywords("a😀b", Some(0.5)), "a😀ab"); + } + + #[test] + fn replace_handles_text_whose_lowercase_is_longer() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("istanbul", Some("IST")); + assert_eq!(kp.replace_keywords("İstanbul istanbul", None), "İstanbul IST"); + } + + #[test] + fn non_ascii_keywords_can_be_listed_looked_up_and_removed() { + let mut kp = KeywordProcessor::new(); + kp.add_keyword("naïve", Some("naive")); + kp.add_keyword("café", None); + let all = kp.get_all_keywords(); + assert_eq!(all.get("naïve"), Some(&"naive".to_string())); + assert_eq!(all.get("café"), Some(&"café".to_string())); + assert_eq!(kp.get_keyword("NAÏVE"), Some("naive".to_string())); + assert!(kp.remove_keyword("naïve")); + assert!(!kp.remove_keyword("naïve")); + assert_eq!(kp.len(), 1); + } } diff --git a/src/lib.rs b/src/lib.rs index c61a95d..a1f6d81 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -107,6 +107,8 @@ //! mod char_set; +mod compiled_trie; +mod fuzzy; mod keyword_processor; mod trie; mod utils; @@ -117,7 +119,10 @@ pub use keyword_processor::KeywordProcessor; use bincode::{deserialize, serialize}; use pyo3::prelude::*; -use pyo3::types::PyBytes; +use pyo3::pybacked::PyBackedStr; +use pyo3::exceptions::PyValueError; +use pyo3::types::{PyBytes, PyString}; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -126,10 +131,8 @@ use std::collections::HashMap; struct PyKeywordProcessor(KeywordProcessor); #[pyclass(module = "blitztext")] -#[derive(Serialize, Deserialize, Clone)] struct PyKeywordMatch { - #[pyo3(get, set)] - pub keyword: String, + keyword: Py, #[pyo3(get, set)] pub similarity: f32, #[pyo3(get, set)] @@ -138,48 +141,89 @@ struct PyKeywordMatch { pub end: usize, } +/// Pickled form of a `PyKeywordMatch`. +#[derive(Serialize, Deserialize)] +struct KeywordMatchState { + keyword: String, + similarity: f32, + start: usize, + end: usize, +} + #[pymethods] impl PyKeywordMatch { #[new] - fn new(keyword: String, similarity: f32, start: usize, end: usize) -> Self { + fn new(py: Python<'_>, keyword: String, similarity: f32, start: usize, end: usize) -> Self { PyKeywordMatch { - keyword, + keyword: PyString::new(py, &keyword).unbind(), similarity, start, end, } } - fn __repr__(&self) -> String { - format!( + #[getter] + fn keyword(&self, py: Python<'_>) -> Py { + self.keyword.clone_ref(py) + } + + #[setter] + fn set_keyword(&mut self, py: Python<'_>, keyword: String) { + self.keyword = PyString::new(py, &keyword).unbind(); + } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( "KeywordMatch(keyword='{}', similarity={:.2}, start={}, end={})", - self.keyword, self.similarity, self.start, self.end - ) + self.keyword.bind(py).to_cow()?, + self.similarity, + self.start, + self.end + )) } - pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> { - *self = deserialize(state.as_bytes()).unwrap(); + pub fn __setstate__(&mut self, py: Python<'_>, state: Bound<'_, PyBytes>) -> PyResult<()> { + let state: KeywordMatchState = deserialize(state.as_bytes()) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + *self = PyKeywordMatch::new(py, state.keyword, state.similarity, state.start, state.end); Ok(()) } pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { - Ok(PyBytes::new_bound(py, &serialize(&self).unwrap())) + let state = KeywordMatchState { + keyword: self.keyword.bind(py).to_cow()?.into_owned(), + similarity: self.similarity, + start: self.start, + end: self.end, + }; + Ok(PyBytes::new(py, &serialize(&state).unwrap())) } - fn __getnewargs__(&self) -> PyResult<(String, f32, usize, usize)> { - Ok((self.keyword.clone(), self.similarity, self.start, self.end)) + fn __getnewargs__(&self, py: Python<'_>) -> (Py, f32, usize, usize) { + (self.keyword.clone_ref(py), self.similarity, self.start, self.end) } } -impl From> for PyKeywordMatch { - fn from(km: KeywordMatch) -> Self { - PyKeywordMatch { - keyword: km.keyword.to_string(), - similarity: km.similarity, - start: km.start, - end: km.end, - } - } +/// Converts matches to Python objects, sharing one string per distinct keyword. +fn to_python_matches( + py: Python<'_>, + matches: Vec>, + keywords: &mut FxHashMap<(*const u8, usize), Py>, +) -> Vec { + matches + .into_iter() + .map(|keyword_match| { + let keyword = keywords + .entry((keyword_match.keyword.as_ptr(), keyword_match.keyword.len())) + .or_insert_with(|| PyString::new(py, keyword_match.keyword).unbind()); + PyKeywordMatch { + keyword: keyword.clone_ref(py), + similarity: keyword_match.similarity, + start: keyword_match.start, + end: keyword_match.end, + } + }) + .collect() } #[pymethods] @@ -205,40 +249,46 @@ impl PyKeywordProcessor { } #[pyo3(signature = (text, threshold=None))] - fn extract_keywords(&self, text: &str, threshold: Option) -> Vec { - self.0 - .extract_keywords(text, threshold) - .into_iter() - .map(PyKeywordMatch::from) - .collect() + fn extract_keywords( + &self, + py: Python<'_>, + text: &str, + threshold: Option, + ) -> PyResult> { + let matches = self.0.extract_keywords(text, checked(threshold)?); + Ok(to_python_matches(py, matches, &mut FxHashMap::default())) } #[pyo3(signature = (texts, threshold=None))] fn parallel_extract_keywords_from_texts( &self, - texts: Vec, + py: Python<'_>, + texts: Vec, threshold: Option, - ) -> Vec> { - self.0 - .parallel_extract_keywords_from_texts(&texts, threshold) + ) -> PyResult>> { + let threshold = checked(threshold)?; + let mut keywords = FxHashMap::default(); + Ok(py + .detach(|| self.0.parallel_extract_keywords_from_texts(&texts, threshold)) .into_iter() - .map(|matches| matches.into_iter().map(PyKeywordMatch::from).collect()) - .collect() + .map(|matches| to_python_matches(py, matches, &mut keywords)) + .collect()) } #[pyo3(signature = (text, threshold=None))] - fn replace_keywords(&self, text: &str, threshold: Option) -> String { - self.0.replace_keywords(text, threshold) + fn replace_keywords(&self, text: &str, threshold: Option) -> PyResult { + Ok(self.0.replace_keywords(text, checked(threshold)?)) } #[pyo3(signature = (texts, threshold=None))] fn parallel_replace_keywords_from_texts( &self, - texts: Vec, + py: Python<'_>, + texts: Vec, threshold: Option, - ) -> Vec { - self.0 - .parallel_replace_keywords_from_texts(&texts, threshold) + ) -> PyResult> { + let threshold = checked(threshold)?; + Ok(py.detach(|| self.0.parallel_replace_keywords_from_texts(&texts, threshold))) } fn get_all_keywords(&self) -> HashMap { @@ -263,11 +313,20 @@ impl PyKeywordProcessor { // https://users.rust-lang.org/t/rust-extension-module-breaks-multiprocessing-no-pickle-support/112402 pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> { - *self = deserialize(state.as_bytes()).unwrap(); + *self = deserialize(state.as_bytes()).map_err(|error| PyValueError::new_err(error.to_string()))?; Ok(()) } pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { - Ok(PyBytes::new_bound(py, &serialize(&self).unwrap())) + Ok(PyBytes::new(py, &serialize(&self).unwrap())) + } +} + +fn checked(threshold: Option) -> PyResult> { + match threshold { + Some(threshold) if !(0.0..=1.0).contains(&threshold) => { + Err(PyValueError::new_err("Threshold must be between 0 and 1")) + } + threshold => Ok(threshold), } } diff --git a/src/trie.rs b/src/trie.rs index 00aea0f..c8f0b1d 100644 --- a/src/trie.rs +++ b/src/trie.rs @@ -1,3 +1,4 @@ +use crate::fuzzy::TrieView; use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; @@ -34,4 +35,60 @@ impl TrieNode { keyword: None, } } + + pub fn count_nodes(&self) -> usize { + let mut count = 0; + let mut pending = vec![self]; + while let Some(node) = pending.pop() { + count += 1; + pending.extend(node.children.values().map(|child| &**child)); + } + count + } +} + +impl<'t> TrieView<'t> for &'t TrieNode { + type Node = &'t TrieNode; + + fn root(&self) -> Self::Node { + self + } + + fn node_id(&self, node: Self::Node) -> usize { + node as *const TrieNode as usize + } + + fn clean_name(&self, node: Self::Node) -> Option<&'t str> { + node.keyword.as_deref() + } + + fn byte_child(&self, node: Self::Node, byte: u8) -> Option { + node.children.get(&(byte as char)).map(|child| &**child) + } + + fn for_each_byte_child(&self, node: Self::Node, mut visit: impl FnMut(u8, Self::Node)) { + for (&key, child) in &node.children { + visit(key as u32 as u8, child); + } + } + + fn chars_to_deepest_keyword(&self, _node: Self::Node) -> Option { + // Every leaf ends a keyword, but the depth is not tracked here. + Some(usize::MAX) + } + + fn longest_prefix(&self, text: &[u8]) -> Option<(&'t str, usize)> { + let mut current_node = *self; + let mut found = None; + for (i, &byte) in text.iter().enumerate() { + match current_node.children.get(&(byte as char)) { + Some(next_node) => current_node = next_node, + None => break, + } + if let Some(keyword) = ¤t_node.keyword { + found = Some((keyword.as_str(), i + 1)); + } + } + found + } } diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 72458b4..13a8ffe 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -42,9 +42,9 @@ fn test_keyword_processor_integration() { // Test fuzzy matching let fuzzy_text = "I love programming in Rast and Pythn. Sometimes I use Jva too."; processor.remove_keyword("go"); - let fuzzy_matches = processor.extract_keywords(fuzzy_text, Some(0.8)); + let fuzzy_matches = processor.extract_keywords(fuzzy_text, Some(0.7)); assert_eq!(fuzzy_matches.len(), 3); - assert!(fuzzy_matches.iter().all(|m| m.similarity >= 0.8)); + assert!(fuzzy_matches.iter().all(|m| m.similarity >= 0.7)); // Test case sensitivity let mut case_sensitive_processor = KeywordProcessor::with_options(true, false); diff --git a/tests/test_python_api.py b/tests/test_python_api.py new file mode 100644 index 0000000..50cacae --- /dev/null +++ b/tests/test_python_api.py @@ -0,0 +1,39 @@ +import math +import pickle + +import pytest +from blitztext import KeywordProcessor +from blitztext import PyKeywordMatch as KeywordMatch + + +@pytest.fixture +def processor(): + kp = KeywordProcessor() + kp.add_keyword("rust", "Rust Lang") + return kp + + +@pytest.mark.parametrize("threshold", [-0.1, 1.5, math.nan]) +def test_threshold_outside_zero_to_one_raises_value_error(processor, threshold): + with pytest.raises(ValueError, match="between 0 and 1"): + processor.extract_keywords("rust", threshold) + with pytest.raises(ValueError): + processor.parallel_replace_keywords_from_texts(["rust"], threshold) + + +def test_corrupt_pickle_raises_value_error(processor): + with pytest.raises(ValueError): + processor.__setstate__(b"\x01\x02") + + +def test_processor_and_matches_survive_pickling(processor): + restored = pickle.loads(pickle.dumps(processor)) + [match] = restored.extract_keywords("I love rust") + assert (match.keyword, match.start, match.end) == ("Rust Lang", 7, 11) + assert repr(pickle.loads(pickle.dumps(match))) == repr(match) + + +def test_match_keyword_can_be_reassigned(): + match = KeywordMatch("a", 1.0, 0, 1) + match.keyword = "b" + assert match.keyword == "b"