refactor(av1): reduce complexity in cdef and loop restoration (S3776)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/av1/src/cdef.rs

@@ -45,7 +45,7 @@ fn cdef_block(fs: &FrameState, out: &mut FrameBuf, r: usize, c: usize, idx: usiz
     let var_str = if (var >> 6) != 0 { (31 - ((var >> 6) as u32).leading_zeros()).min(12) as i32 } else { 0 };
     pri = if var != 0 { (pri * (4 + var_str) + 8) >> 4 } else { 0 };
     let damping = cp.damping as i32 + coeff_shift as i32;
-    filter(fs, out, 0, r, c, pri, sec, damping, dir);
+    filter(fs, out, 0, (r, c), (pri, sec), damping, dir);
     if fs.num_planes == 1 {
         return true;
     }
@@ -53,8 +53,8 @@ fn cdef_block(fs: &FrameState, out: &mut FrameBuf, r: usize, c: usize, idx: usiz
     let sec = (cp.uv_sec[idx] << coeff_shift) as i32;
     let dir = if pri == 0 { 0 } else { CDEF_UV_DIR[fs.ssx][fs.ssy][y_dir] as usize };
     let damping = cp.damping as i32 + coeff_shift as i32 - 1;
-    filter(fs, out, 1, r, c, pri, sec, damping, dir);
-    filter(fs, out, 2, r, c, pri, sec, damping, dir);
+    filter(fs, out, 1, (r, c), (pri, sec), damping, dir);
+    filter(fs, out, 2, (r, c), (pri, sec), damping, dir);
     true
 }
 
@@ -124,33 +124,74 @@ fn constrain(diff: i32, threshold: i32, adj: u32) -> i32 {
     if diff < 0 { -mag } else { mag }
 }
 
-#[allow(clippy::too_many_arguments)]
-fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, r: usize, c: usize, pri: i32, sec: i32, damping: i32, dir: usize) {
-    if pri == 0 && sec == 0 {
-        // Every tap contributes 0 and the clamp to [min, max] keeps x: CdefFrame = CurrFrame.
-        return;
+/// Tap offsets, weights and strengths of one block.
+struct Taps {
+    pri: i32,
+    sec: i32,
+    pri_adj: u32,
+    sec_adj: u32,
+    pri_off: [isize; 2],
+    sec_off: [[isize; 2]; 2],
+    pri_taps: [i32; 2],
+    sec_taps: [i32; 2],
+}
+
+impl Taps {
+    fn new(coeff_shift: u32, (pri, sec): (i32, i32), damping: i32, dir: usize) -> Taps {
+        let pri_tap = ((pri >> coeff_shift) & 1) as usize;
+        let off = |d: usize, k: usize| CDEF_DIRECTIONS[d][k][0] as isize * PAD_STRIDE as isize + CDEF_DIRECTIONS[d][k][1] as isize;
+        let pri_off = [off(dir, 0), off(dir, 1)];
+        let d_lo = (dir + 6) & 7;
+        let d_hi = (dir + 2) & 7;
+        let sec_off = [[off(d_lo, 0), off(d_hi, 0)], [off(d_lo, 1), off(d_hi, 1)]];
+        let pri_taps = [CDEF_PRI_TAPS[pri_tap][0] as i32, CDEF_PRI_TAPS[pri_tap][1] as i32];
+        let sec_taps = [CDEF_SEC_TAPS[pri_tap][0] as i32, CDEF_SEC_TAPS[pri_tap][1] as i32];
+        let adj = |t: i32| if t == 0 { 0 } else { (damping - (31 - (t as u32).leading_zeros()) as i32).max(0) as u32 };
+        Taps { pri, sec, pri_adj: adj(pri), sec_adj: adj(sec), pri_off, sec_off, pri_taps, sec_taps }
     }
-    let coeff_shift = fs.bit_depth - 8;
+
+    /// The 12 (offset, weight, strength, adjustment) taps in filtering order.
+    fn tap_list(&self) -> [(isize, i32, i32, u32); 12] {
+        let mut taps = [(0isize, 0i32, 0i32, 0u32); 12];
+        let mut n = 0;
+        for k in 0..2 {
+            for sign in [-1isize, 1] {
+                taps[n] = (sign * self.pri_off[k], self.pri_taps[k], self.pri, self.pri_adj);
+                taps[n + 1] = (sign * self.sec_off[k][0], self.sec_taps[k], self.sec, self.sec_adj);
+                taps[n + 2] = (sign * self.sec_off[k][1], self.sec_taps[k], self.sec, self.sec_adj);
+                n += 3;
+            }
+        }
+        taps
+    }
+}
+
+/// Position and size of a block within its plane.
+#[derive(Clone, Copy)]
+struct BlockRect {
+    x0: usize,
+    y0: usize,
+    w: usize,
+    h: usize,
+}
+
+/// The block with two samples of padding on every side, plus whether every sample is available
+/// (CdefAvailable: its 4x4 block is inside the frame).
+fn padded_block(fs: &FrameState, plane: usize, src: &Plane, rect: BlockRect) -> ([i32; PAD_STRIDE * PAD_STRIDE], bool) {
     let (sub_x, sub_y) = if plane > 0 { (fs.ssx, fs.ssy) } else { (0, 0) };
-    let x0 = (c * 4) >> sub_x;
-    let y0 = (r * 4) >> sub_y;
-    let w = 8 >> sub_x;
-    let h = 8 >> sub_y;
-    let src = &fs.cur.planes[plane];
     let (mi_rows, mi_cols) = (fs.fh.mi_rows as isize, fs.fh.mi_cols as isize);
-    // CdefAvailable per sample: its 4x4 block is inside the frame.
     let mut buf = [CDEF_NA; PAD_STRIDE * PAD_STRIDE];
     let mut all_avail = true;
-    for i in 0..h + 4 {
-        let y = (y0 + i) as isize - 2;
+    for i in 0..rect.h + 4 {
+        let y = (rect.y0 + i) as isize - 2;
         let cand_r = (y << sub_y) >> 2;
         if y < 0 || cand_r >= mi_rows {
             all_avail = false;
             continue;
         }
         let row = src.row(y as usize);
-        for j in 0..w + 4 {
-            let x = (x0 + j) as isize - 2;
+        for j in 0..rect.w + 4 {
+            let x = (rect.x0 + j) as isize - 2;
             let cand_c = (x << sub_x) >> 2;
             if x >= 0 && cand_c < mi_cols {
                 buf[i * PAD_STRIDE + j] = row[x as usize] as i32;
@@ -159,70 +200,70 @@ fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, r: usize, c: usize,
             }
         }
     }
-    let pri_tap = ((pri >> coeff_shift) & 1) as usize;
-    let off = |d: usize, k: usize| CDEF_DIRECTIONS[d][k][0] as isize * PAD_STRIDE as isize + CDEF_DIRECTIONS[d][k][1] as isize;
-    let pri_off = [off(dir, 0), off(dir, 1)];
-    let d_lo = (dir + 6) & 7;
-    let d_hi = (dir + 2) & 7;
-    let sec_off = [[off(d_lo, 0), off(d_hi, 0)], [off(d_lo, 1), off(d_hi, 1)]];
-    let pri_taps = [CDEF_PRI_TAPS[pri_tap][0] as i32, CDEF_PRI_TAPS[pri_tap][1] as i32];
-    let sec_taps = [CDEF_SEC_TAPS[pri_tap][0] as i32, CDEF_SEC_TAPS[pri_tap][1] as i32];
-    let adj = |t: i32| if t == 0 { 0 } else { (damping - (31 - (t as u32).leading_zeros()) as i32).max(0) as u32 };
-    let (pri_adj, sec_adj) = (adj(pri), adj(sec));
+    (buf, all_avail)
+}
+
+/// Filters one sample whose taps may be unavailable (CDEF_NA).
+fn filter_sample(buf: &[i32; PAD_STRIDE * PAD_STRIDE], taps: &Taps, ci: isize, x: i32) -> u16 {
+    let mut sum = 0i32;
+    let mut max = x;
+    let mut min = x;
+    for k in 0..2 {
+        for sign in [-1isize, 1] {
+            let p = buf[(ci + sign * taps.pri_off[k]) as usize];
+            if p != CDEF_NA {
+                if taps.pri != 0 {
+                    sum += taps.pri_taps[k] * constrain(p - x, taps.pri, taps.pri_adj);
+                }
+                max = max.max(p);
+                min = min.min(p);
+            }
+            for s in 0..2 {
+                let q = buf[(ci + sign * taps.sec_off[k][s]) as usize];
+                if q != CDEF_NA {
+                    if taps.sec != 0 {
+                        sum += taps.sec_taps[k] * constrain(q - x, taps.sec, taps.sec_adj);
+                    }
+                    max = max.max(q);
+                    min = min.min(q);
+                }
+            }
+        }
+    }
+    (x + ((8 + sum - (sum < 0) as i32) >> 4)).clamp(min, max) as u16
+}
+
+fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, (r, c): (usize, usize), strengths: (i32, i32), damping: i32, dir: usize) {
+    let (pri, sec) = strengths;
+    if pri == 0 && sec == 0 {
+        // Every tap contributes 0 and the clamp to [min, max] keeps x: CdefFrame = CurrFrame.
+        return;
+    }
+    let coeff_shift = fs.bit_depth - 8;
+    let (sub_x, sub_y) = if plane > 0 { (fs.ssx, fs.ssy) } else { (0, 0) };
+    let rect = BlockRect { x0: (c * 4) >> sub_x, y0: (r * 4) >> sub_y, w: 8 >> sub_x, h: 8 >> sub_y };
+    let src = &fs.cur.planes[plane];
+    let (buf, all_avail) = padded_block(fs, plane, src, rect);
+    let taps = Taps::new(coeff_shift, strengths, damping, dir);
     let dst = &mut out.planes[plane];
     if all_avail {
         // Every tap is available: whole rows at a time (constrain() with threshold 0 is 0).
-        let mut taps = [(0isize, 0i32, 0i32, 0u32); 12];
-        let mut n = 0;
-        for k in 0..2 {
-            for sign in [-1isize, 1] {
-                taps[n] = (sign * pri_off[k], pri_taps[k], pri, pri_adj);
-                taps[n + 1] = (sign * sec_off[k][0], sec_taps[k], sec, sec_adj);
-                taps[n + 2] = (sign * sec_off[k][1], sec_taps[k], sec, sec_adj);
-                n += 3;
-            }
-        }
-        let rows = (y0, x0, h);
-        match w {
-            8 => filter_rows::<8>(&buf, &taps, dst, rows),
-            _ => filter_rows::<4>(&buf, &taps, dst, rows),
+        let list = taps.tap_list();
+        let rows = (rect.y0, rect.x0, rect.h);
+        match rect.w {
+            8 => filter_rows::<8>(&buf, &list, dst, rows),
+            _ => filter_rows::<4>(&buf, &list, dst, rows),
         }
         return;
     }
-    for i in 0..h {
+    for i in 0..rect.h {
         // The centre sample is read even where CdefAvailable is 0 (8x8 blocks straddling an odd
         // MiRows / MiCols edge).
-        let srow = &src.row(y0 + i)[x0..x0 + w];
-        let drow = &mut dst.row_mut(y0 + i)[x0..x0 + w];
-        for j in 0..w {
+        let srow = &src.row(rect.y0 + i)[rect.x0..rect.x0 + rect.w];
+        let drow = &mut dst.row_mut(rect.y0 + i)[rect.x0..rect.x0 + rect.w];
+        for j in 0..rect.w {
             let ci = ((i + 2) * PAD_STRIDE + j + 2) as isize;
-            let x = srow[j] as i32;
-            let mut sum = 0i32;
-            let mut max = x;
-            let mut min = x;
-            for k in 0..2 {
-                for sign in [-1isize, 1] {
-                    let p = buf[(ci + sign * pri_off[k]) as usize];
-                    if p != CDEF_NA {
-                        if pri != 0 {
-                            sum += pri_taps[k] * constrain(p - x, pri, pri_adj);
-                        }
-                        max = max.max(p);
-                        min = min.min(p);
-                    }
-                    for s in 0..2 {
-                        let q = buf[(ci + sign * sec_off[k][s]) as usize];
-                        if q != CDEF_NA {
-                            if sec != 0 {
-                                sum += sec_taps[k] * constrain(q - x, sec, sec_adj);
-                            }
-                            max = max.max(q);
-                            min = min.min(q);
-                        }
-                    }
-                }
-            }
-            drow[j] = (x + ((8 + sum - (sum < 0) as i32) >> 4)).clamp(min, max) as u16;
+            drow[j] = filter_sample(&buf, &taps, ci, srow[j] as i32);
         }
     }
 }

crates/av1/src/restoration.rs

@@ -73,57 +73,70 @@ impl LrCtx<'_> {
     }
 }
 
+/// Position and size of a (at most 4x4) restoration block within its plane.
+#[derive(Clone, Copy)]
+struct LrBlock {
+    x: i32,
+    y: i32,
+    w: i32,
+    h: i32,
+}
+
 /// Loop restoration (7.17): returns LrFrame given UpscaledCurrFrame and UpscaledCdefFrame.
 pub(crate) fn loop_restoration(fs: &FrameState, cur: &FrameBuf, cdef: &FrameBuf) -> FrameBuf {
-    let fh = &fs.fh;
     let mut lr = cdef.clone();
-    let up_w = fh.upscaled_width as i32;
-    let fr_h = fh.frame_height as i32;
-    let bd = fs.bit_depth;
     for plane in 0..fs.num_planes {
-        let rtype_frame = fh.lr.frame_restoration_type[plane];
-        if rtype_frame == RESTORE_NONE as u8 {
+        if fs.fh.lr.frame_restoration_type[plane] == RESTORE_NONE as u8 {
             continue;
         }
-        let (sub_x, sub_y) = if plane > 0 { (fs.ssx as i32, fs.ssy as i32) } else { (0, 0) };
-        let unit_size = fh.lr.loop_restoration_size[plane] as i32;
-        let unit_rows = fs.lr_unit_rows[plane] as i32;
-        let unit_cols = fs.lr_unit_cols[plane] as i32;
-        let plane_end_x = ((up_w + sub_x) >> sub_x) - 1;
-        let plane_end_y = ((fr_h + sub_y) >> sub_y) - 1;
-        let mut row = 0;
-        while row * 4 < fr_h {
-            let luma_y = row * 4;
-            let stripe_num = (luma_y + 8) / 64;
-            let stripe_start_y = (-8 + stripe_num * 64) >> sub_y;
-            let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1;
-            let ctx = LrCtx { cur: &cur.planes[plane], cdef: &cdef.planes[plane], plane_end_x, plane_end_y, stripe_start_y, stripe_end_y };
-            let unit_row = (unit_rows - 1).min(((row * 4 + 8) >> sub_y) / unit_size);
-            let mut col = 0;
-            while col * 4 < up_w {
-                let unit_col = (unit_cols - 1).min(((col * 4) >> sub_x) / unit_size);
-                let x = (col * 4) >> sub_x;
-                let y = (row * 4) >> sub_y;
-                let w = (4 >> sub_x).min(plane_end_x - x + 1);
-                let h = (4 >> sub_y).min(plane_end_y - y + 1);
-                if w > 0 && h > 0 {
-                    let ui = (unit_row * unit_cols + unit_col) as usize;
-                    match fs.lr.lr_type[plane][ui] {
-                        t if t == RESTORE_WIENER as u8 => wiener(fs, &ctx, &mut lr.planes[plane], plane, ui, x, y, w, h, bd),
-                        t if t == RESTORE_SGRPROJ as u8 => self_guided(fs, &ctx, &mut lr.planes[plane], plane, ui, x, y, w, h, bd),
-                        _ => {}
-                    }
-                }
-                col += 1;
-            }
-            row += 1;
-        }
+        restore_plane(fs, cur, cdef, &mut lr.planes[plane], plane);
     }
     lr
 }
 
-#[allow(clippy::too_many_arguments)]
-fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize, x: i32, y: i32, w: i32, h: i32, bd: u32) {
+fn restore_plane(fs: &FrameState, cur: &FrameBuf, cdef: &FrameBuf, out: &mut Plane, plane: usize) {
+    let fh = &fs.fh;
+    let up_w = fh.upscaled_width as i32;
+    let fr_h = fh.frame_height as i32;
+    let bd = fs.bit_depth;
+    let (sub_x, sub_y) = if plane > 0 { (fs.ssx as i32, fs.ssy as i32) } else { (0, 0) };
+    let unit_size = fh.lr.loop_restoration_size[plane] as i32;
+    let unit_rows = fs.lr_unit_rows[plane] as i32;
+    let unit_cols = fs.lr_unit_cols[plane] as i32;
+    let plane_end_x = ((up_w + sub_x) >> sub_x) - 1;
+    let plane_end_y = ((fr_h + sub_y) >> sub_y) - 1;
+    let mut row = 0;
+    while row * 4 < fr_h {
+        let luma_y = row * 4;
+        let stripe_num = (luma_y + 8) / 64;
+        let stripe_start_y = (-8 + stripe_num * 64) >> sub_y;
+        let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1;
+        let ctx = LrCtx { cur: &cur.planes[plane], cdef: &cdef.planes[plane], plane_end_x, plane_end_y, stripe_start_y, stripe_end_y };
+        let unit_row = (unit_rows - 1).min(((row * 4 + 8) >> sub_y) / unit_size);
+        let mut col = 0;
+        while col * 4 < up_w {
+            let unit_col = (unit_cols - 1).min(((col * 4) >> sub_x) / unit_size);
+            let x = (col * 4) >> sub_x;
+            let y = (row * 4) >> sub_y;
+            let w = (4 >> sub_x).min(plane_end_x - x + 1);
+            let h = (4 >> sub_y).min(plane_end_y - y + 1);
+            if w > 0 && h > 0 {
+                let ui = (unit_row * unit_cols + unit_col) as usize;
+                let blk = LrBlock { x, y, w, h };
+                match fs.lr.lr_type[plane][ui] {
+                    t if t == RESTORE_WIENER as u8 => wiener(fs, &ctx, out, (plane, ui), blk, bd),
+                    t if t == RESTORE_SGRPROJ as u8 => self_guided(fs, &ctx, out, (plane, ui), blk, bd),
+                    _ => {}
+                }
+            }
+            col += 1;
+        }
+        row += 1;
+    }
+}
+
+fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, (plane, ui): (usize, usize), blk: LrBlock, bd: u32) {
+    let LrBlock { x, y, w, h } = blk;
     let round0: u32 = if bd == 12 { 5 } else { 3 };
     let round1: u32 = if bd == 12 { 9 } else { 11 };
     let coef = &fs.lr.lr_wiener[plane][ui];
@@ -166,13 +179,13 @@ fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize
     }
 }
 
-#[allow(clippy::too_many_arguments)]
-fn self_guided(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize, x: i32, y: i32, w: i32, h: i32, bd: u32) {
+fn self_guided(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, (plane, ui): (usize, usize), blk: LrBlock, bd: u32) {
+    let LrBlock { x, y, w, h } = blk;
     let set = fs.lr.lr_sgr_set[plane][ui] as usize;
     let mut flt0 = [[0i32; 4]; 4];
     let mut flt1 = [[0i32; 4]; 4];
-    box_filter(ctx, x, y, w, h, set, 0, bd, &mut flt0);
-    box_filter(ctx, x, y, w, h, set, 1, bd, &mut flt1);
+    box_filter(ctx, blk, (set, 0), bd, &mut flt0);
+    box_filter(ctx, blk, (set, 1), bd, &mut flt1);
     let xqd = fs.lr.lr_sgr_xqd[plane][ui];
     let w0 = xqd[0] as i32;
     let w1 = xqd[1] as i32;
@@ -198,8 +211,56 @@ fn round2_64(x: i64, n: u32) -> i64 {
     if n == 0 { x } else { (x + (1i64 << (n - 1))) >> n }
 }
 
-#[allow(clippy::too_many_arguments)]
-fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usize, bd: u32, f: &mut [[i32; 4]; 4]) {
+/// Box sums of the samples within `r` of (x, y): (sum of squares, sum).
+fn box_sums(ctx: &LrCtx, x: i32, y: i32, r: i32) -> (i64, i64) {
+    let mut a = 0i64;
+    let mut b = 0i64;
+    for dy in -r..=r {
+        for dx in -r..=r {
+            let c = ctx.sample(x + dx, y + dy) as i64;
+            a += c * c;
+            b += c;
+        }
+    }
+    (a, b)
+}
+
+/// Self-guided filter constants of one pass.
+struct SgrConsts {
+    n: i64,
+    s: i64,
+    one_over_n: i64,
+}
+
+/// (A, B) coefficients of one sample from its box sums.
+fn sgr_coefficients(a: i64, b: i64, k: &SgrConsts, bd: u32) -> (i64, i64) {
+    let a = round2_64(a, 2 * (bd - 8));
+    let d = round2_64(b, bd - 8);
+    let p = (a * k.n - d * d).max(0);
+    let z = round2_64(p * k.s, SGRPROJ_MTABLE_BITS as u32);
+    let a2 = if z >= 255 {
+        256
+    } else if z == 0 {
+        1
+    } else {
+        ((z << SGRPROJ_SGR_BITS) + z / 2) / (z + 1)
+    };
+    let b2 = ((1i64 << SGRPROJ_SGR_BITS) - a2) * b * k.one_over_n;
+    (a2, round2_64(b2, SGRPROJ_RECIP_BITS as u32))
+}
+
+fn box_weight(pass: usize, i: i32, dx: i32, dy: i32) -> i64 {
+    if pass == 0 {
+        if ((i + dy) & 1) != 0 { if dx == 0 { 6 } else { 5 } } else { 0 }
+    } else if dx == 0 || dy == 0 {
+        4
+    } else {
+        3
+    }
+}
+
+fn box_filter(ctx: &LrCtx, blk: LrBlock, (set, pass): (usize, usize), bd: u32, f: &mut [[i32; 4]; 4]) {
+    let LrBlock { x, y, w, h } = blk;
     let r = SGR_PARAMS[set][pass * 2] as i32;
     if r == 0 {
         return;
@@ -207,35 +268,15 @@ fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usi
     let eps = SGR_PARAMS[set][pass * 2 + 1] as i64;
     let n = ((2 * r + 1) * (2 * r + 1)) as i64;
     let n2e = n * n * eps;
-    let s = ((1i64 << SGRPROJ_MTABLE_BITS) + n2e / 2) / n2e;
-    let one_over_n = ((1i64 << SGRPROJ_RECIP_BITS) + n / 2) / n;
+    let k = SgrConsts { n, s: ((1i64 << SGRPROJ_MTABLE_BITS) + n2e / 2) / n2e, one_over_n: ((1i64 << SGRPROJ_RECIP_BITS) + n / 2) / n };
     let mut a_arr = [[0i64; 6]; 6];
     let mut b_arr = [[0i64; 6]; 6];
     for i in -1..h + 1 {
         for j in -1..w + 1 {
-            let mut a = 0i64;
-            let mut b = 0i64;
-            for dy in -r..=r {
-                for dx in -r..=r {
-                    let c = ctx.sample(x + j + dx, y + i + dy) as i64;
-                    a += c * c;
-                    b += c;
-                }
-            }
-            let a = round2_64(a, 2 * (bd - 8));
-            let d = round2_64(b, bd - 8);
-            let p = (a * n - d * d).max(0);
-            let z = round2_64(p * s, SGRPROJ_MTABLE_BITS as u32);
-            let a2 = if z >= 255 {
-                256
-            } else if z == 0 {
-                1
-            } else {
-                ((z << SGRPROJ_SGR_BITS) + z / 2) / (z + 1)
-            };
-            let b2 = ((1i64 << SGRPROJ_SGR_BITS) - a2) * b * one_over_n;
+            let (a, b) = box_sums(ctx, x + j, y + i, r);
+            let (a2, b2) = sgr_coefficients(a, b, &k, bd);
             a_arr[(i + 1) as usize][(j + 1) as usize] = a2;
-            b_arr[(i + 1) as usize][(j + 1) as usize] = round2_64(b2, SGRPROJ_RECIP_BITS as u32);
+            b_arr[(i + 1) as usize][(j + 1) as usize] = b2;
         }
     }
     for i in 0..h {
@@ -245,13 +286,7 @@ fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usi
             let mut b = 0i64;
             for dy in -1..=1i32 {
                 for dx in -1..=1i32 {
-                    let weight = if pass == 0 {
-                        if ((i + dy) & 1) != 0 { if dx == 0 { 6 } else { 5 } } else { 0 }
-                    } else if dx == 0 || dy == 0 {
-                        4
-                    } else {
-                        3
-                    };
+                    let weight = box_weight(pass, i, dx, dy);
                     a += weight * a_arr[(i + dy + 1) as usize][(j + dx + 1) as usize];
                     b += weight * b_arr[(i + dy + 1) as usize][(j + dx + 1) as usize];
                 }