Skip to content

Commit

Permalink
fn rav1d_prepare_intra_edges: clean up and make safe (#690)
Browse files Browse the repository at this point in the history
  • Loading branch information
@kkysen
kkysen authored Feb 6, 2024
2 parents 223b536 + ddddc08 commit 2df91e1
Show file tree
Hide file tree
Showing 2 changed files with 253 additions and 218 deletions.
252 changes: 103 additions & 149 deletions src/ipred_prepare.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,15 @@ use crate::src::levels::SMOOTH_H_PRED;
use crate::src::levels::SMOOTH_PRED;
use crate::src::levels::SMOOTH_V_PRED;
use crate::src::levels::TOP_DC_PRED;
use crate::src::levels::VERT_LEFT_PRED;
use crate::src::levels::VERT_PRED;
use crate::src::levels::Z1_PRED;
use crate::src::levels::Z2_PRED;
use crate::src::levels::Z3_PRED;
use bitflags::bitflags;
use libc::memcpy;
use libc::ptrdiff_t;
use std::cmp;
use std::ffi::c_int;
use std::ffi::c_uint;
use std::ffi::c_void;
use std::slice;

#[inline]
pub fn sm_flag(b: &BlockContext, idx: usize) -> c_int {
Expand Down Expand Up @@ -120,59 +117,61 @@ static av1_intra_prediction_edges: [av1_intra_prediction_edge; N_IMPL_INTRA_PRED
b
};

pub unsafe fn rav1d_prepare_intra_edges<BD: BitDepth>(
pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
x: c_int,
have_left: c_int,
have_left: bool,
y: c_int,
have_top: c_int,
have_top: bool,
w: c_int,
h: c_int,
edge_flags: EdgeFlags,
dst: *const BD::Pixel,
dst: &[BD::Pixel], // contains 4*h first rows of picture, last row in slice contains 4*w samples
stride: ptrdiff_t,
prefilter_toplevel_sb_edge: *const BD::Pixel,
prefilter_toplevel_sb_edge: Option<&[BD::Pixel]>,
mut mode: IntraPredMode,
angle: *mut c_int,
angle: &mut c_int,
tw: c_int,
th: c_int,
filter_edge: c_int,
topleft_out: *mut BD::Pixel,
topleft_out: &mut [BD::Pixel],
topleft_out_offset: usize, // position of top-left sample in `topleft_out`
bd: BD,
) -> IntraPredMode {
assert!(y < h && x < w);

let bitdepth = bd.bitdepth();
if !(y < h && x < w) {
unreachable!();
}
match mode as c_uint {
1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 => {
*angle = av1_mode_to_angle_map
[(mode as c_uint).wrapping_sub(VERT_PRED as c_int as c_uint) as usize]
as c_int
+ 3 * *angle;
let stride = BD::pxstride(stride as usize) as isize;

let dst_offset = 4 * x as usize
+ (if stride >= 0 { 4 * y } else { 4 * (h - y) - 1 }) as usize * stride.unsigned_abs();

match mode {
VERT_PRED..=VERT_LEFT_PRED => {
*angle = av1_mode_to_angle_map[(mode - VERT_PRED) as usize] as c_int + 3 * *angle;
if *angle <= 90 {
mode = (if *angle < 90 && have_top != 0 {
Z1_PRED as c_int
mode = if *angle < 90 && have_top {
Z1_PRED
} else {
VERT_PRED as c_int
}) as IntraPredMode;
VERT_PRED
};
} else if *angle < 180 {
mode = Z2_PRED;
} else {
mode = (if *angle > 180 && have_left != 0 {
Z3_PRED as c_int
mode = if *angle > 180 && have_left {
Z3_PRED
} else {
HOR_PRED as c_int
}) as IntraPredMode;
HOR_PRED
};
}
}
0 | 12 => {
mode = av1_mode_conv[mode as usize][have_left as usize][have_top as usize]
as IntraPredMode;
DC_PRED | PAETH_PRED => {
mode = av1_mode_conv[mode as usize][have_left as usize][have_top as usize];
}
_ => {}
}
let mut dst_top: *const BD::Pixel = 0 as *const BD::Pixel;
if have_top != 0

// `dst_top` starts with either the top or top-left sample depending on whether have_left is true
let dst_top = if have_top
&& (av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::TOP)
Expand All @@ -182,177 +181,132 @@ pub unsafe fn rav1d_prepare_intra_edges<BD: BitDepth>(
|| av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::LEFT)
&& have_left == 0)
&& !have_left)
{
if !prefilter_toplevel_sb_edge.is_null() {
dst_top = &*prefilter_toplevel_sb_edge.offset((x * 4) as isize) as *const BD::Pixel;
let px_have = cmp::min(8 * tw, 4 * (w - x)) as usize;
let n = px_have + have_left as usize;
if let Some(prefilter_toplevel_sb_edge) = prefilter_toplevel_sb_edge {
let offset = (x * 4) as usize - have_left as usize;
&prefilter_toplevel_sb_edge[offset..][..n]
} else {
dst_top = &*dst.offset(-(BD::pxstride(stride as usize) as isize)) as *const BD::Pixel;
&dst[(dst_offset as isize - stride) as usize - have_left as usize..][..n]
}
}
} else {
&[]
};

if av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::LEFT)
{
let sz = th << 2;
let left: *mut BD::Pixel = &mut *topleft_out.offset(-sz as isize) as *mut BD::Pixel;
if have_left != 0 {
let px_have = cmp::min(sz, h - y << 2);
let mut i = 0;
while i < px_have {
*left.offset((sz - 1 - i) as isize) =
*dst.offset(BD::pxstride(stride as usize) as isize * i as isize - 1);
i += 1;
let sz = 4 * th as usize;
let left = &mut topleft_out[topleft_out_offset - sz..];
if have_left {
let px_have = cmp::min(sz, (h - y << 2) as usize);
for i in 0..px_have {
left[sz - 1 - i] = dst[i * stride as usize + dst_offset - 1];
}
if px_have < sz {
BD::pixel_set(
slice::from_raw_parts_mut(left, (sz - px_have).try_into().unwrap()),
*left.offset((sz - px_have) as isize),
(sz - px_have).try_into().unwrap(),
);
BD::pixel_set(left, left[sz - px_have], sz - px_have);
}
} else {
BD::pixel_set(
slice::from_raw_parts_mut(left, sz.try_into().unwrap()),
if have_top != 0 {
*dst_top
left,
if have_top {
dst_top[0] // have_left is always false
} else {
((1 << bitdepth >> 1) + 1).as_::<BD::Pixel>()
},
sz.try_into().unwrap(),
sz,
);
}
if av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::BOTTOM_LEFT)
{
let have_bottomleft = (if have_left == 0 || y + th >= h {
0 as c_int as c_uint
let bottom_left = &mut topleft_out[topleft_out_offset - 2 * sz..];
let have_bottomleft = if !have_left || y + th >= h {
false
} else {
edge_flags as c_uint & EDGE_I444_LEFT_HAS_BOTTOM as c_int as c_uint
}) as c_int;
if have_bottomleft != 0 {
let px_have_0 = cmp::min(sz, h - y - th << 2);
let mut i_0 = 0;
while i_0 < px_have_0 {
*left.offset(-(i_0 + 1) as isize) = *dst.offset(
((sz + i_0) as isize * BD::pxstride(stride as usize) as isize - 1 as isize)
as isize,
);
i_0 += 1;
(edge_flags & EDGE_I444_LEFT_HAS_BOTTOM) != 0
};
if have_bottomleft {
let px_have = cmp::min(sz, (h - y - th << 2) as usize);
for i in 0..px_have {
bottom_left[sz - 1 - i] = dst[(sz + i) * stride as usize + dst_offset - 1];
}
if px_have_0 < sz {
BD::pixel_set(
slice::from_raw_parts_mut(
left.offset(-(sz as isize)),
(sz - px_have_0).try_into().unwrap(),
),
*left.offset(-px_have_0 as isize),
(sz - px_have_0).try_into().unwrap(),
);
if px_have < sz {
BD::pixel_set(bottom_left, bottom_left[sz - px_have], sz - px_have);
}
} else {
BD::pixel_set(
slice::from_raw_parts_mut(left.offset(-(sz as isize)), sz.try_into().unwrap()),
*left.offset(0),
sz.try_into().unwrap(),
);
BD::pixel_set(bottom_left, bottom_left[sz], sz);
}
}
}
if av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::TOP)
{
let sz_0 = tw << 2;
let top: *mut BD::Pixel = &mut *topleft_out.offset(1) as *mut BD::Pixel;
if have_top != 0 {
let px_have_1 = cmp::min(sz_0, w - x << 2);
memcpy(
top as *mut c_void,
dst_top as *const c_void,
(px_have_1 << 1) as usize,
);
if px_have_1 < sz_0 {
BD::pixel_set(
slice::from_raw_parts_mut(
top.offset(px_have_1 as isize),
(sz_0 - px_have_1).try_into().unwrap(),
),
*top.offset((px_have_1 - 1) as isize),
(sz_0 - px_have_1).try_into().unwrap(),
);
let sz = 4 * tw as usize;
let top = &mut topleft_out[topleft_out_offset + 1..];
if have_top {
let px_have = cmp::min(sz, (w - x << 2) as usize);
BD::pixel_copy(top, &dst_top[have_left as usize..], px_have);
if px_have < sz {
let fill_value = top[px_have - 1];
BD::pixel_set(&mut top[px_have..], fill_value, sz - px_have);
}
} else {
BD::pixel_set(
slice::from_raw_parts_mut(top, sz_0.try_into().unwrap()),
if have_left != 0 {
*dst.offset(-(1 as c_int) as isize)
top,
if have_left {
dst[dst_offset - 1]
} else {
((1 << bitdepth >> 1) - 1).as_::<BD::Pixel>()
},
sz_0.try_into().unwrap(),
sz,
);
}
if av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::TOP_RIGHT)
{
let have_topright = (if have_top == 0 || x + tw >= w {
0 as c_int as c_uint
let have_topright = if !have_top || x + tw >= w {
false
} else {
edge_flags as c_uint & EDGE_I444_TOP_HAS_RIGHT as c_int as c_uint
}) as c_int;
if have_topright != 0 {
let px_have_2 = cmp::min(sz_0, w - x - tw << 2);
memcpy(
top.offset(sz_0 as isize) as *mut c_void,
&*dst_top.offset(sz_0 as isize) as *const BD::Pixel as *const c_void,
(px_have_2 << 1) as usize,
);
if px_have_2 < sz_0 {
BD::pixel_set(
slice::from_raw_parts_mut(
top.offset(sz_0 as isize).offset(px_have_2 as isize),
(sz_0 - px_have_2).try_into().unwrap(),
),
*top.offset((sz_0 + px_have_2 - 1) as isize),
(sz_0 - px_have_2).try_into().unwrap(),
);
(edge_flags & EDGE_I444_TOP_HAS_RIGHT) != 0
};
if have_topright {
let top_right = &mut top[sz..];
let px_have = cmp::min(sz, (w - x - tw << 2) as usize);
BD::pixel_copy(top_right, &dst_top[sz + have_left as usize..], px_have);
if px_have < sz {
let fill_value = top_right[px_have - 1];
BD::pixel_set(&mut top_right[px_have..], fill_value, sz - px_have);
}
} else {
BD::pixel_set(
slice::from_raw_parts_mut(top.offset(sz_0 as isize), sz_0.try_into().unwrap()),
*top.offset((sz_0 - 1) as isize),
sz_0.try_into().unwrap(),
);
let fill_value = top[sz - 1];
BD::pixel_set(&mut top[sz..], fill_value, sz);
}
}
}
if av1_intra_prediction_edges[mode as usize]
.needs
.contains(Needs::TOP_LEFT)
{
if have_left != 0 {
*topleft_out = (if have_top != 0 {
(*dst_top.offset(-(1 as c_int) as isize)).as_::<c_int>()
} else {
(*dst.offset(-(1 as c_int) as isize)).as_::<c_int>()
})
.as_::<BD::Pixel>();
// top-left sample and immediate neighbours
let corner =
<&mut [_; 3]>::try_from(&mut topleft_out[topleft_out_offset - 1..][..3]).unwrap();
corner[1] = if have_top {
dst_top[0]
} else if have_left {
dst[dst_offset - 1]
} else {
*topleft_out = (if have_top != 0 {
(*dst_top).as_::<c_int>()
} else {
(1 as c_int) << bitdepth >> 1
})
.as_::<BD::Pixel>();
}
if mode as c_uint == Z2_PRED as c_int as c_uint && tw + th >= 6 && filter_edge != 0 {
*topleft_out = (((*topleft_out.offset(-(1 as c_int) as isize)).as_::<c_int>()
+ (*topleft_out.offset(1)).as_::<c_int>())
* 5
+ (*topleft_out.offset(0)).as_::<c_int>() * 6
(1 << bitdepth >> 1).as_::<BD::Pixel>()
};
if mode == Z2_PRED && tw + th >= 6 && filter_edge != 0 {
corner[1] = ((corner[0].as_::<c_int>() + corner[2].as_::<c_int>()) * 5
+ corner[1].as_::<c_int>() * 6
+ 8
>> 4)
.as_::<BD::Pixel>();
Expand Down
Loading

0 comments on commit 2df91e1

Please sign in to comment.