From cacc762810b51742916575360504b84fa81fdca2 Mon Sep 17 00:00:00 2001
From: kageru <kageru@encode.moe>
Date: Sun, 2 Jun 2019 22:44:52 +0200
Subject: [PATCH] Support all common formats

---
 Cargo.toml  |   5 +--
 src/lib.rs  |   6 ++-
 src/mask.rs | 103 ++++++++++++++++++++++++++++++++++++++++++----------
 3 files changed, 89 insertions(+), 25 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 07af6d0..0fdd256 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2018"
 
 [dependencies]
 failure = "0.1"
-faster = { git = "https://github.com/AdamNiederer/faster.git" }
+#faster = { git = "https://github.com/AdamNiederer/faster.git" }
 lazy_static = "1.3.0"
 num = "0.2"
 rand = "0.6"
@@ -17,6 +17,3 @@ vapoursynth-sys = "0.2"
 [lib]
 crate-type = ["cdylib"]
 
-[profile.release]
-debug = true
-
diff --git a/src/lib.rs b/src/lib.rs
index c8dadaf..1f74814 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -111,14 +111,16 @@ make_filter_function! {
             None => 10.0
         };
         if let Property::Constant(format) = clip.info().format {
-            if format.bits_per_sample() == 32 && format.sample_type() == SampleType::Float {
+            if !(format.sample_type() == SampleType::Float && format.bits_per_sample() != 32) {
                 return Ok(Some(Box::new(Mask {
                     source: clip,
                     luma_scaling
                 })));
+            } else {
+                bail!("Half precision float input is not supported");
             }
         }
-        bail!("Currently, only constant 32bit float input is supported" )
+        bail!("Variable format input is not supported")
     }
 }
 
diff --git a/src/mask.rs b/src/mask.rs
index 0b5a17b..23ea7ab 100644
--- a/src/mask.rs
+++ b/src/mask.rs
@@ -1,14 +1,13 @@
 use super::PLUGIN_NAME;
 use failure::Error;
+use num::{NumCast, ToPrimitive};
 use std::fmt::Debug;
-use vapoursynth::api::API;
+use std::ops::Shl;
 use vapoursynth::core::CoreRef;
 use vapoursynth::format::ColorFamily;
-use vapoursynth::frame::{FrameRef, FrameRefMut};
-use vapoursynth::node::Node;
 use vapoursynth::plugins::{Filter, FrameContext};
+use vapoursynth::prelude::*;
 use vapoursynth::video_info::{Property, VideoInfo};
-use std::hint::unreachable_unchecked;
 
 pub struct Mask<'core> {
     pub source: Node<'core>,
@@ -17,10 +16,10 @@ pub struct Mask<'core> {
 
 lazy_static! {
     static ref FLOAT_RANGE: Vec<f32> = {
-        [0f32; 1000]
+        [0f32; 256]
             .iter()
             .enumerate()
-            .map(|(i, _f)| (i as f32) * 0.001)
+            .map(|(i, _f)| (i as f32) * (1.0 / 256.0))
             .collect::<Vec<_>>()
     };
 }
@@ -41,6 +40,51 @@ fn from_property<T: Debug + Clone + Copy + Eq + PartialEq>(prop: Property<T>) ->
     }
 }
 
+fn filter_for_int<T>(
+    frame: &mut FrameRefMut,
+    src_frame: FrameRef,
+    depth: u8,
+    average: f32,
+    luma_scaling: f32,
+) where
+    T: Component + NumCast + Clone + Shl<u8> + Debug,
+    <T as Shl<u8>>::Output: ToPrimitive,
+{
+    let max = ((1 << depth) - 1) as f32;
+    let lut: Vec<T> = FLOAT_RANGE
+        .iter()
+        .map(|x| {
+            T::from(get_mask_value(*x, average, luma_scaling) * max).unwrap()
+        })
+        .collect();
+    for row in 0..frame.height(0) {
+        for (pixel, src_pixel) in frame
+            .plane_row_mut::<T>(0, row)
+            .iter_mut()
+            .zip(src_frame.plane_row::<T>(0, row).iter())
+        {
+            let i = <usize as NumCast>::from(src_pixel.clone()).unwrap() >> (depth - 8);
+            *pixel = lut[i].clone();
+        }
+    }
+}
+
+fn filter_for_float(frame: &mut FrameRefMut, src_frame: FrameRef, average: f32, luma_scaling: f32) {
+    let lut: Vec<f32> = FLOAT_RANGE
+        .iter()
+        .map(|x| get_mask_value(*x, average, luma_scaling))
+        .collect();
+    for row in 0..frame.height(0) {
+        for (pixel, src_pixel) in frame
+            .plane_row_mut::<f32>(0, row)
+            .iter_mut()
+            .zip(src_frame.plane_row::<f32>(0, row).iter())
+        {
+            *pixel = lut[(src_pixel * 255f32) as usize];
+        }
+    }
+}
+
 impl<'core> Filter<'core> for Mask<'core> {
     fn video_info(&self, _api: API, _core: CoreRef<'core>) -> Vec<VideoInfo<'core>> {
         let info = self.source.info();
@@ -99,24 +143,45 @@ impl<'core> Filter<'core> for Mask<'core> {
         })?;
         let average = match src_frame.props().get::<f64>("PlaneStatsAverage") {
             Ok(average) => average as f32,
-            Err(_) => panic!(format!(
+            Err(_) => bail!(format!(
                 "{}: you need to run std.PlaneStats on the clip before calling this function.",
                 PLUGIN_NAME
             )),
         };
 
-        let lut: Vec<f32> = FLOAT_RANGE
-            .iter()
-            .map(|x| get_mask_value(*x, average, self.luma_scaling))
-            .collect();
-
-        for row in 0..frame.height(0) {
-            for (pixel, src_pixel) in frame
-                .plane_row_mut::<f32>(0, row)
-                .iter_mut()
-                .zip(src_frame.plane_row::<f32>(0, row).iter())
-            {
-                *pixel = lut[(src_pixel * 1000f32) as usize];
+        match from_property(self.source.info().format).sample_type() {
+            SampleType::Integer => {
+                let depth = from_property(self.source.info().format).bits_per_sample();
+                match depth {
+                    _ if depth <= 8 => filter_for_int::<u8>(
+                        &mut frame,
+                        src_frame,
+                        depth,
+                        average,
+                        self.luma_scaling,
+                    ),
+                    _ if 8 < depth && depth <= 16 => filter_for_int::<u16>(
+                        &mut frame,
+                        src_frame,
+                        depth,
+                        average,
+                        self.luma_scaling,
+                    ),
+                    _ if 16 < depth && depth <= 32 => filter_for_int::<u32>(
+                        &mut frame,
+                        src_frame,
+                        depth,
+                        average,
+                        self.luma_scaling,
+                    ),
+                    _ => bail!(format!(
+                        "{}: input depth {} not supported",
+                        PLUGIN_NAME, depth
+                    )),
+                }
+            }
+            SampleType::Float => {
+                filter_for_float(&mut frame, src_frame, average, self.luma_scaling);
             }
         }
         Ok(frame.into())