♻️ refactor: rustify writer utility test_free_rows

Author: SigureMo

packages/biliass/rust/src/lib.rs

@@ -31,5 +31,7 @@ fn biliass_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(python::py_convert_color, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_get_zoom_factor, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_convert_flash_rotation, m)?)?;
+    m.add_function(wrap_pyfunction!(python::py_test_free_rows, m)?)?;
+    m.add_class::<python::PyRows>()?;
     Ok(())
 }

packages/biliass/rust/src/python/comment.rs

@@ -27,7 +27,7 @@ impl PyCommentPosition {
 
 #[pyclass(name = "Comment", frozen)]
 pub struct PyComment {
-    inner: comment::Comment,
+    pub inner: comment::Comment,
 }
 
 impl PyComment {
@@ -96,7 +96,7 @@ impl PyComment {
 
 #[pyclass(name = "OptionComment", frozen)]
 pub struct PyOptionComment {
-    inner: Option<comment::Comment>,
+    pub inner: Option<comment::Comment>,
 }
 
 impl PyOptionComment {
@@ -107,27 +107,27 @@ impl PyOptionComment {
 
 #[pymethods]
 impl PyOptionComment {
-    fn is_none(&self) -> bool {
+    pub fn is_none(&self) -> bool {
         self.inner.is_none()
     }
 
-    fn is_some(&self) -> bool {
+    pub fn is_some(&self) -> bool {
         self.inner.is_some()
     }
 
-    fn unwrap(&self) -> PyComment {
+    pub fn unwrap(&self) -> PyComment {
         PyComment::new(self.inner.clone().unwrap())
     }
 
     #[staticmethod]
-    fn from_comment(comment: &PyComment) -> Self {
+    pub fn from_comment(comment: &PyComment) -> Self {
         PyOptionComment {
             inner: Some(comment.inner.clone()),
         }
     }
 
     #[staticmethod]
-    fn none() -> Self {
+    pub fn none() -> Self {
         PyOptionComment { inner: None }
     }
 

packages/biliass/rust/src/python/mod.rs

@@ -8,5 +8,5 @@ pub use proto::{PyDanmakuElem, PyDmSegMobileReply};
 pub use reader::{py_read_comments_from_protobuf, py_read_comments_from_xml};
 pub use writer::{
     py_ass_escape, py_convert_color, py_convert_flash_rotation, py_convert_timestamp,
-    py_get_zoom_factor,
+    py_get_zoom_factor, py_test_free_rows, PyRows,
 };

packages/biliass/rust/src/python/writer.rs

@@ -1,7 +1,10 @@
-use crate::writer;
+use crate::python;
+use crate::writer::{self, rows};
 
 use pyo3::prelude::*;
 
+use super::PyOptionComment;
+
 #[pyfunction(name = "convert_timestamp")]
 pub fn py_convert_timestamp(timestamp: f64) -> PyResult<String> {
     Ok(writer::utils::convert_timestamp(timestamp))
@@ -40,7 +43,55 @@ pub fn py_convert_flash_rotation(
     ))
 }
 
-// #[pyfunction(name = "test_rows")]
-// pub fn py_test_rows(rows: Vec<Vec<PyRef<crate::python::PyOptionComment>>>) -> PyResult<()> {
-//     Ok(())
-// }
+#[pyclass(name = "Rows")]
+pub struct PyRows {
+    pub inner: rows::Rows,
+}
+
+#[pymethods]
+impl PyRows {
+    #[new]
+    fn new(num_types: usize, capacity: usize) -> Self {
+        let mut rows: rows::Rows = Vec::new();
+        for _ in 0..num_types {
+            let mut type_rows = Vec::with_capacity(capacity);
+            for _ in 0..capacity {
+                type_rows.push(None);
+            }
+            rows.push(type_rows);
+        }
+        PyRows { inner: rows }
+    }
+
+    fn get(&self, row: usize, col: usize) -> PyOptionComment {
+        PyOptionComment::new(self.inner[row][col].clone())
+    }
+
+    fn set(&mut self, row: usize, col: usize, value: PyRef<PyOptionComment>) {
+        self.inner[row][col] = value.inner.clone();
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+#[pyfunction(name = "test_free_rows")]
+pub fn py_test_free_rows(
+    rows: &python::writer::PyRows,
+    comment: &crate::python::PyComment,
+    row: usize,
+    width: u32,
+    height: u32,
+    bottom_reserved: u32,
+    duration_marquee: f64,
+    duration_still: f64,
+) -> PyResult<usize> {
+    Ok(writer::rows::test_free_rows(
+        &rows.inner,
+        &comment.inner,
+        row,
+        width,
+        height,
+        bottom_reserved,
+        duration_marquee,
+        duration_still,
+    ))
+}

packages/biliass/rust/src/writer/mod.rs

@@ -1 +1,2 @@
+pub mod rows;
 pub mod utils;

packages/biliass/rust/src/writer/rows.rs

@@ -0,0 +1,57 @@
+use crate::comment::{Comment, CommentPosition};
+
+pub type Rows = Vec<Vec<Option<Comment>>>;
+
+#[allow(clippy::too_many_arguments)]
+pub fn test_free_rows(
+    rows: &Rows,
+    comment: &Comment,
+    row: usize,
+    width: u32,
+    height: u32,
+    bottom_reserved: u32,
+    duration_marquee: f64,
+    duration_still: f64,
+) -> usize {
+    let mut res = 0;
+    let rowmax = (height - bottom_reserved) as usize;
+    let mut target_row = None;
+    let comment_pos_id = comment.pos.clone() as usize;
+    if comment.pos == CommentPosition::Bottom || comment.pos == CommentPosition::Top {
+        let mut current_row = row;
+        while current_row < rowmax && (res as f32) < comment.height {
+            if target_row != rows[comment_pos_id][current_row] {
+                target_row = rows[comment_pos_id][current_row].clone();
+                if let Some(target_row) = target_row.clone() {
+                    if target_row.timeline + duration_still > comment.timeline {
+                        break;
+                    }
+                }
+            }
+            current_row += 1;
+            res += 1;
+        }
+    } else {
+        let threshold_time: f64 = comment.timeline
+            - duration_marquee * (1.0 - width as f64 / (comment.width as f64 + width as f64));
+        let mut current_row = row;
+        while current_row < rowmax && (res as f32) < comment.height {
+            if target_row != rows[comment_pos_id][current_row] {
+                target_row = rows[comment_pos_id][current_row].clone();
+                if let Some(target_row) = target_row.clone() {
+                    if target_row.timeline > threshold_time
+                        || target_row.timeline
+                            + target_row.width as f64 * duration_marquee
+                                / (target_row.width as f64 + width as f64)
+                            > comment.timeline
+                    {
+                        break;
+                    }
+                }
+            }
+            current_row += 1;
+            res += 1;
+        }
+    }
+    res
+}

packages/biliass/src/biliass/_core.pyi

@@ -73,3 +73,18 @@ def get_zoom_factor(source_size: tuple[int, int], target_size: tuple[int, int])
 def convert_flash_rotation(
     rot_y: float, rot_z: float, x: float, y: float, width: float, height: float
 ) -> tuple[float, float, float, float, float, float, float]: ...
+def test_free_rows(
+    rows: Rows,
+    comment: Comment,
+    row: int,
+    width: int,
+    height: int,
+    bottom_reserved: int,
+    duration_marquee: float,
+    duration_still: float,
+) -> int: ...
+
+class Rows:
+    def __init__(self, num_types: int, capacity: int) -> None: ...
+    def get(self, row: int, col: int) -> OptionComment: ...
+    def set(self, row: int, col: int, OptionComment) -> None: ...

packages/biliass/src/biliass/biliass.py

@@ -13,21 +13,22 @@
     Comment,
     CommentPosition,
     OptionComment,
+    Rows,
     ass_escape,
     convert_color,
     convert_flash_rotation,
     convert_timestamp,
     get_zoom_factor,
     read_comments_from_protobuf,
     read_comments_from_xml,
+    test_free_rows,
 )
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
 
 T = TypeVar("T")
-Rows = list[list[OptionComment]]
 
 
 def read_comments_bilibili_xml(text: str | bytes, fontsize: float) -> list[Comment]:
@@ -245,7 +246,7 @@ def process_comments(
     styleid = f"biliass_{random.randint(0, 0xFFFF):04x}"
     ass = AssText()
     ass.write_head(width, height, fontface, fontsize, alpha, styleid)
-    rows = [[OptionComment.none()] * (height - bottom_reserved + 1) for i in range(4)]
+    rows = Rows(4, height - bottom_reserved + 1)
     for idx, comment in enumerate(comments):
         if progress_callback and idx % 1000 == 0:
             progress_callback(idx, len(comments))
@@ -315,58 +316,13 @@ def process_comments(
     return ass.to_string()
 
 
-def test_free_rows(
-    rows: Rows,
-    comment: Comment,
-    row: int,
-    width: int,
-    height: int,
-    bottom_reserved: int,
-    duration_marquee: float,
-    duration_still: float,
-) -> int:
-    res = 0
-    rowmax = height - bottom_reserved
-    target_row = OptionComment.none()
-    comment_pos_id = comment.pos.id
-    if comment.pos in (CommentPosition.Bottom, CommentPosition.Top):
-        while row < rowmax and res < comment.height:
-            if target_row != rows[comment_pos_id][row]:
-                target_row = rows[comment_pos_id][row]
-                if target_row.is_some() and target_row.unwrap().timeline + duration_still > comment.timeline:
-                    break
-            row += 1
-            res += 1
-    else:
-        try:
-            threshold_time = comment.timeline - duration_marquee * (1 - width / (comment.width + width))
-        except ZeroDivisionError:
-            threshold_time = comment.timeline - duration_marquee
-        while row < rowmax and res < comment.height:
-            if target_row != rows[comment_pos_id][row]:
-                target_row = rows[comment_pos_id][row]
-                try:
-                    if target_row.is_some() and (
-                        target_row.unwrap().timeline > threshold_time
-                        or target_row.unwrap().timeline
-                        + target_row.unwrap().width * duration_marquee / (target_row.unwrap().width + width)
-                        > comment.timeline
-                    ):
-                        break
-                except ZeroDivisionError:
-                    pass
-            row += 1
-            res += 1
-    return res
-
-
 def find_alternative_row(rows: Rows, comment: Comment, height, bottom_reserved):
     res = 0
     comment_pos_id = comment.pos.id
     for row in range(height - bottom_reserved - math.ceil(comment.height)):
-        if rows[comment_pos_id][row].is_none():
+        if rows.get(comment_pos_id, row).is_none():
             return row
-        elif rows[comment_pos_id][row].unwrap().timeline < rows[comment_pos_id][res].unwrap().timeline:
+        elif rows.get(comment_pos_id, row).unwrap().timeline < rows.get(comment_pos_id, res).unwrap().timeline:
             res = row
     return res
 
@@ -375,7 +331,7 @@ def mark_comment_row(rows: Rows, comment: Comment, row: int):
     comment_pos_id = comment.pos.id
     try:
         for i in range(row, row + math.ceil(comment.height)):
-            rows[comment_pos_id][i] = OptionComment.from_comment(comment)
+            rows.set(comment_pos_id, i, OptionComment.from_comment(comment))
     except IndexError:
         pass