v0.1.1

Author: swoiow

.github/workflows/release.yml

@@ -50,6 +50,10 @@ jobs:
         python-version:
           - "3.13"
 
+    env:
+      PYTHONIOENCODING: "utf-8"
+      PYTHONUTF8: "1"
+
     steps:
       - name: Checkout
         uses: actions/checkout@v4

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "poptrie"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 
 [lib]

build_bin.py

@@ -3,7 +3,7 @@
 
 
 class BinBuilder:
-    NODE_SIZE = 68  # 32字节Child_BM + 32字节Leaf_BM + 4字节Base_Offset
+    NODE_SIZE = 72  # 32字节Child_BM + 32字节Leaf_BM + 4字节Base_Offset + 4字节Padding
 
     def __init__(self):
         # 根节点：{ 'children': {byte: node_dict}, 'is_leaf': bool }
@@ -111,12 +111,12 @@ def save(self, output_path):
 
                 if nodes_with_children:
                     # 记录跳转到下一层这些子节点的起始偏移
-                    final_data.extend(struct.pack("<I", next_layer_start_offset))
+                    final_data.extend(struct.pack("<II", next_layer_start_offset, 0))
                     next_layer.extend(nodes_with_children)
                     next_layer_start_offset += len(nodes_with_children) * self.NODE_SIZE
                 else:
                     # 没有子节点可跳转
-                    final_data.extend(struct.pack("<I", 0))
+                    final_data.extend(struct.pack("<II", 0, 0))
 
             current_layer = next_layer
 

src/lib.rs

@@ -1,8 +1,8 @@
 use memmap2::Mmap;
 use pyo3::prelude::*;
-use pyo3::types::PyBytes;
 use rayon::prelude::*;
 use std::fs::File;
+use std::net::IpAddr;
 
 #[pyclass]
 struct IpSearcher {
@@ -15,113 +15,111 @@ impl IpSearcher {
     fn new(path: String) -> PyResult<Self> {
         let file = File::open(path)?;
         let mmap = unsafe { Mmap::map(&file)? };
+        if mmap.len() % Self::NODE_SIZE != 0 {
+            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                "Invalid bin file: alignment mismatch (expected 72).",
+            ));
+        }
         Ok(IpSearcher { mmap })
     }
 
     /// 核心查询逻辑：支持 IPv4 (4字节) 和 IPv6 (16字节)
     fn is_china_ip(&self, ip_bytes: &[u8]) -> bool {
-        let mut curr_ptr: usize = 0;
-        let node_size: usize = 68;
+        let mut cursor: usize = 0;
 
         for &byte in ip_bytes {
             // 节点布局: 0-31 ChildBitmap, 32-63 LeafBitmap, 64-67 BaseOffset
-            let child_bm = &self.mmap[curr_ptr..curr_ptr + 32];
-            let leaf_bm = &self.mmap[curr_ptr + 32..curr_ptr + 64];
+            let child_bitmap = &self.mmap[cursor..cursor + 32];
+            let leaf_bitmap = &self.mmap[cursor + 32..cursor + 64];
 
             // 1. 检查当前步长是否匹配 (Leaf)
-            if self.check_bit(leaf_bm, byte) {
+            if self.check_bit(leaf_bitmap, byte) {
                 return true;
             }
 
             // 2. 检查是否有子节点
-            if !self.check_bit(child_bm, byte) {
+            if !self.check_bit(child_bitmap, byte) {
                 return false;
             }
 
             // 3. 计算跳转偏移 (Popcount)
             // 读取 4 字节的 BaseOffset (小端序)
             let base_offset =
-                u32::from_le_bytes(self.mmap[curr_ptr + 64..curr_ptr + 68].try_into().unwrap())
+                u32::from_le_bytes(self.mmap[cursor + 64..cursor + 68].try_into().unwrap())
                     as usize;
 
             // 获取当前字节之前的 '1' 的数量，确定子节点索引
-            let index = self.get_popcount(child_bm, byte);
-            curr_ptr = base_offset + (index * node_size);
+            let index = self.get_popcount(child_bitmap, byte);
+            cursor = base_offset + (index * Self::NODE_SIZE);
         }
         false
     }
 
-    // 这里的 Bound<'_, PyBytes> 允许我们直接访问 Python 的内存
-    fn batch_check(&self, ip_list: Vec<Bound<'_, PyBytes>>) -> Vec<bool> {
-        ip_list
-            .into_iter()
-            .map(|py_bytes| {
-                // as_bytes() 返回 &[u8]，不需要拷贝数据
-                self.is_china_ip(py_bytes.as_bytes())
-            })
-            .collect()
-    }
-
-    /// 极致性能版：接收一个扁平化的字节流（每 4 或 16 字节代表一个 IP）
+    /// 极致性能版：接收扁平化字节流（每 4 或 16 字节代表一个 IP）
     fn batch_check_packed(&self, packed_ips: &[u8], is_v6: bool) -> Vec<bool> {
-        let stride = if is_v6 { 16 } else { 4 };
+        let ip_stride = if is_v6 { 16 } else { 4 };
 
         // 使用 chunks_exact 确保每次切出固定长度的 IP 字节块
         // 这是极致性能的关键：内存完全连续，没有 Python 对象开销
         packed_ips
-            .chunks_exact(stride)
+            .chunks_exact(ip_stride)
             .map(|ip_chunk| self.is_china_ip(ip_chunk))
             .collect()
     }
 
-    fn batch_check_packed_parallel(
-        &self,
-        py: Python<'_>,
-        packed_ips: &[u8],
-        is_v6: bool,
-    ) -> Vec<bool> {
-        let stride = if is_v6 { 16 } else { 4 };
-
-        // par_chunks_exact 是 Rayon 提供的并行切片方法
+    fn batch_check_strings(&self, py: Python<'_>, ips: Vec<String>) -> Vec<bool> {
         py.allow_threads(|| {
-            packed_ips
-                .par_chunks_exact(stride)
-                .map(|ip_chunk| self.is_china_ip(ip_chunk))
+            ips.into_par_iter()
+                .map(|ip_str| match ip_str.parse::<IpAddr>() {
+                    Ok(IpAddr::V4(v4)) => self.is_china_ip(&v4.octets()),
+                    Ok(IpAddr::V6(v6)) => self.is_china_ip(&v6.octets()),
+                    Err(_) => false,
+                })
                 .collect()
         })
     }
 }
 
 impl IpSearcher {
+    const NODE_SIZE: usize = 72;
+
     #[inline]
     fn check_bit(&self, bitmap: &[u8], byte: u8) -> bool {
-        let idx = (byte >> 3) as usize;
-        let bit = 7 - (byte % 8); // 对应 Python 的 (1 << (7 - (k % 8)))
-        (bitmap[idx] >> bit) & 1 == 1
+        let byte_index = (byte >> 3) as usize;
+        let bit_index = 7 - (byte % 8); // 对应 Python 的 (1 << (7 - (k % 8)))
+        (bitmap[byte_index] >> bit_index) & 1 == 1
     }
 
     #[inline]
     fn get_popcount(&self, bitmap: &[u8], byte: u8) -> usize {
-        let byte_idx = (byte >> 3) as usize;
-        let bit_in_byte = 7 - (byte % 8);
-        let mut count = 0;
-
-        // 1. 累加之前所有字节中 1 的个数
-        for i in 0..byte_idx {
+        let byte_index = (byte >> 3) as usize;
+        let bit_index = 7 - (byte % 8);
+        let mut count: usize = 0;
+
+        // 1. 以 u64 为单位统计，减少循环次数
+        let full_byte_count = byte_index;
+        let chunk_count = full_byte_count / 8;
+        for i in 0..chunk_count {
+            let start = i * 8;
+            let value =
+                u64::from_le_bytes(bitmap[start..start + 8].try_into().unwrap());
+            count += value.count_ones() as usize;
+        }
+        for i in (chunk_count * 8)..full_byte_count {
             count += bitmap[i].count_ones() as usize;
         }
 
         // 2. 累加当前字节中，目标位“左侧”所有 1 的个数
         // 我们需要一个掩码来保留比 bit_in_byte 更高的位
         // 例如：如果 bit_in_byte 是 5 (二进制 00100000)，
         // 我们需要掩码 11000000 来计算它之前的 1
-        let mask = if bit_in_byte == 7 {
+        let mask = if bit_index == 7 {
             0
         } else {
-            0xFF << (bit_in_byte + 1)
+            0xFF << (bit_index + 1)
         };
 
-        count += (bitmap[byte_idx] & mask).count_ones() as usize;
+        count += (bitmap[byte_index] & mask).count_ones() as usize;
 
         // 返回值即为该子节点在子节点数组中的索引（从 0 开始）
         count

tests/test_poptrie.py

@@ -55,11 +55,7 @@ def test_ipv6_boundary(self):
 
     def test_batch_check(self):
         ips = ["1.0.1.1", "8.8.8.8", "240e::1", "2001:db8::"]
-        ip_bytes_list = [
-            socket.inet_pton(socket.AF_INET6 if ':' in ip else socket.AF_INET, ip)
-            for ip in ips
-        ]
-        results = self.searcher.batch_check(ip_bytes_list)
+        results = self.searcher.batch_check_strings(ips)
         self.assertEqual(results, [True, False, True, False])
 
     def test_batch_check_packed(self):