nipdf/
function.rs

1use crate::{
2    ObjectValueError, Result,
3    file::ObjectResolver,
4    object::{
5        FromSchemaContainer, Object, ObjectWithResolver, PdfObject as _, RootPdfObject as _,
6        RuntimeObjectId,
7    },
8};
9use educe::Educe;
10#[cfg(test)]
11use mockall::automock;
12use nipdf_macro::{TryFromIntObject, pdf_object};
13use num_traits::ToPrimitive;
14use prescript::{PdfFunc, sname};
15use snafu::{OptionExt as _, ResultExt as _, ensure_whatever};
16use tinyvec::{TinyVec, tiny_vec};
17
18#[derive(Debug, Clone, Copy, PartialEq)]
19pub struct Domain<T = f32> {
20    pub start: T,
21    pub end: T,
22}
23
24impl<T: PartialOrd + Copy> Domain<T> {
25    pub fn new(start: T, end: T) -> Self {
26        Self { start, end }
27    }
28
29    pub fn clamp(&self, x: T) -> T {
30        num_traits::clamp(x, self.start, self.end)
31    }
32
33    pub fn is_zero(&self) -> bool {
34        self.start == self.end
35    }
36}
37
38/// Default domain is [0, 1]
39pub fn default_domain() -> Domain {
40    Domain::new(0.0, 1.0)
41}
42
43impl TryFrom<ObjectWithResolver<'_, '_>> for Domain<f32> {
44    type Error = ObjectValueError;
45
46    fn try_from(obj: ObjectWithResolver<'_, '_>) -> Result<Self, Self::Error> {
47        let arr = obj.into_schema_array()?;
48        ensure_whatever!(
49            arr.len() == 2,
50            "expected array with 2 elements, but got {}",
51            arr.len()
52        );
53        Ok(Self::new(
54            arr.required_object(0)?.number()?,
55            arr.required_object(1)?.number()?,
56        ))
57    }
58}
59
60impl TryFrom<ObjectWithResolver<'_, '_>> for Domain<u32> {
61    type Error = ObjectValueError;
62
63    fn try_from(obj: ObjectWithResolver<'_, '_>) -> Result<Self, Self::Error> {
64        let arr = obj.into_schema_array()?;
65        ensure_whatever!(
66            arr.len() == 2,
67            "expected array with 2 elements, but got {}",
68            arr.len()
69        );
70        Ok(Self::new(
71            arr.required_object(0)?.int()? as u32,
72            arr.required_object(1)?.int()? as u32,
73        ))
74    }
75}
76
77#[derive(Debug, PartialEq, Clone, Educe)]
78#[educe(Deref)]
79pub struct Domains<T = f32>(pub Vec<Domain<T>>);
80
81impl TryFrom<&Object> for Domains<f32> {
82    type Error = ObjectValueError;
83
84    fn try_from(obj: &Object) -> Result<Self, Self::Error> {
85        let arr = obj.as_arr()?;
86        ensure_whatever!(arr.len() % 2 == 0, "even number of elements expected");
87        let mut domains = Vec::with_capacity(arr.len() / 2);
88        arr.chunks_exact(2)
89            .map(|chunk| {
90                Ok::<_, ObjectValueError>(Domain::new(chunk[0].number()?, chunk[1].number()?))
91            })
92            .collect::<Result<Vec<_>, _>>()?
93            .into_iter()
94            .for_each(|domain| domains.push(domain));
95        Ok(Self(domains))
96    }
97}
98
99impl TryFrom<ObjectWithResolver<'_, '_>> for Domains<f32> {
100    type Error = ObjectValueError;
101
102    fn try_from(obj: ObjectWithResolver<'_, '_>) -> Result<Self, Self::Error> {
103        let arr = obj.into_schema_array()?;
104        ensure_whatever!(arr.len() % 2 == 0, "even number of elements expected");
105        let mut domains = Vec::with_capacity(arr.len() / 2);
106        for i in (0..arr.len()).step_by(2) {
107            domains.push(Domain::new(
108                arr.required_object(i)?.number()?,
109                arr.required_object(i + 1)?.number()?,
110            ));
111        }
112        Ok(Self(domains))
113    }
114}
115
116impl TryFrom<&Object> for Domains<u32> {
117    type Error = ObjectValueError;
118
119    fn try_from(obj: &Object) -> Result<Self, Self::Error> {
120        let arr = obj.as_arr()?;
121        let mut domains = Vec::with_capacity(arr.len() / 2);
122        ensure_whatever!(arr.len() % 2 == 0, "even number of elements expected");
123        arr.chunks_exact(2)
124            .map(|chunk| {
125                Ok::<_, ObjectValueError>(Domain::new(
126                    chunk[0].int()? as u32,
127                    chunk[1].int()? as u32,
128                ))
129            })
130            .collect::<Result<Vec<_>, _>>()?
131            .into_iter()
132            .for_each(|domain| domains.push(domain));
133        Ok(Self(domains))
134    }
135}
136impl Domains {
137    /// Function input argument count
138    pub fn n(&self) -> usize {
139        self.0.len()
140    }
141}
142
143pub type FunctionValue = TinyVec<[f32; 4]>;
144
145trait InnerFunction {
146    type Signature: Signature + std::fmt::Debug;
147
148    fn do_call(&self, args: &[f32]) -> Result<FunctionValue> {
149        let args = self.signature().clip_args(args);
150        let r = self.inner_call(args)?;
151        for v in &r {
152            ensure_whatever!(!v.is_nan(), "NaN value");
153        }
154        Ok(self.signature().clip_returns(r))
155    }
156
157    fn signature(&self) -> &Self::Signature;
158
159    /// Called by `self.call()`, args and return value are clipped by signature.
160    fn inner_call(&self, args: FunctionValue) -> Result<FunctionValue>;
161
162    /// Return the domain stops of the function. Helps to build shading stops.
163    fn stops(&self) -> impl Iterator<Item = f32> + 'static;
164}
165
166#[cfg_attr(test, automock)]
167pub trait Function {
168    fn call(&self, args: &[f32]) -> Result<FunctionValue>;
169    fn stops(&self) -> Box<dyn Iterator<Item = f32>>;
170}
171
172impl<Inner: InnerFunction> Function for Inner {
173    fn call(&self, args: &[f32]) -> Result<FunctionValue> {
174        self.do_call(args)
175    }
176
177    fn stops(&self) -> Box<dyn Iterator<Item = f32>> {
178        Box::new(self.stops())
179    }
180}
181
182impl Function for Box<dyn Function> {
183    fn call(&self, args: &[f32]) -> Result<FunctionValue> {
184        self.as_ref().call(args)
185    }
186
187    fn stops(&self) -> Box<dyn Iterator<Item = f32>> {
188        self.as_ref().stops()
189    }
190}
191
192impl<'a, 'b> FromSchemaContainer<'a, 'b> for Box<dyn Function> {
193    fn create(o: &'b Object, r: &'b ObjectResolver<'a>) -> Result<Self, ObjectValueError> {
194        let id: Option<RuntimeObjectId> = o.reference().ok().map(Into::into);
195        let dict = r.resolve_reference(o)?.as_dict()?;
196
197        // Get function type from dictionary
198        let function_type = dict
199            .get(&sname("FunctionType"))
200            .ok_or(ObjectValueError::DictKeyNotFound)?;
201        let function_type = Type::try_from(function_type)?;
202
203        // Create function based on type
204        match function_type {
205            Type::Sampled => {
206                let dict = SampledFunctionDict::new(
207                    id.whatever_context::<_, ObjectValueError>(
208                        "Sampled function should be root object",
209                    )?,
210                    dict,
211                    r,
212                )?;
213                Ok(Box::new(
214                    dict.func()
215                        .whatever_context::<_, ObjectValueError>("Parse Sampled function")?,
216                ))
217            }
218
219            Type::ExponentialInterpolation => {
220                let dict = ExponentialInterpolationFunctionDict::new(dict, r)?;
221                Ok(Box::new(
222                    dict.func().whatever_context::<_, ObjectValueError>(
223                        "Parse Exponential Interpolation function",
224                    )?,
225                ))
226            }
227
228            Type::Stitching => {
229                let dict = StitchingFunctionDict::new(dict, r)?;
230                Ok(Box::new(
231                    dict.func()
232                        .whatever_context::<_, ObjectValueError>("Parse Stitching function")?,
233                ))
234            }
235
236            Type::PostScriptCalculator => {
237                // Get required fields
238                let domain = dict
239                    .get(&sname("Domain"))
240                    .ok_or(ObjectValueError::DictKeyNotFound)
241                    .and_then(Domains::try_from)?;
242
243                let range = dict
244                    .get(&sname("Range"))
245                    .ok_or(ObjectValueError::DictKeyNotFound)
246                    .and_then(Domains::try_from)?;
247
248                let signature = Type04Signature::new(domain, range);
249
250                // Get stream data
251                let stream = r.resolve_reference(o)?.as_stream()?;
252                let script = stream
253                    .decode(r)
254                    .whatever_context::<_, ObjectValueError>("decode stream")?;
255
256                Ok(Box::new(PostScriptFunction::new(
257                    signature,
258                    script.into_owned().into_boxed_slice(),
259                )))
260            }
261        }
262    }
263}
264
265#[derive(Debug, Clone, Copy, PartialEq, TryFromIntObject)]
266pub enum Type {
267    Sampled = 0,
268    ExponentialInterpolation = 2,
269    Stitching = 3,
270    PostScriptCalculator = 4,
271}
272
273pub struct PostScriptFunction {
274    signature: Type04Signature,
275    f: PdfFunc,
276}
277
278impl PostScriptFunction {
279    pub fn new(signature: Type04Signature, script: Box<[u8]>) -> Self {
280        Self {
281            f: PdfFunc::new(script, signature.n_returns()),
282            signature,
283        }
284    }
285}
286
287impl InnerFunction for PostScriptFunction {
288    type Signature = Type04Signature;
289
290    fn signature(&self) -> &Self::Signature {
291        &self.signature
292    }
293
294    #[doc = " Called by `self.call()`, args and return value are clipped by signature."]
295    fn inner_call(&self, args: FunctionValue) -> Result<FunctionValue> {
296        let args = args.into_iter().collect::<Vec<_>>();
297        let r = self
298            .f
299            .exec(&args)
300            .whatever_context::<_, ObjectValueError>("exec function")?;
301        Ok(r.into_iter().collect())
302    }
303
304    fn stops(&self) -> impl Iterator<Item = f32> + 'static {
305        Err("TODO: PostScriptFunction::stops()").into_iter()
306    }
307}
308
309pub trait Signature {
310    fn clip_args(&self, args: &[f32]) -> TinyVec<[f32; 4]>;
311    fn clip_returns(&self, returns: FunctionValue) -> FunctionValue;
312}
313
314/// Function signature for Type 2 and 3, clip input args and returns.
315#[derive(Debug, PartialEq, Clone)]
316pub struct Type23Signature {
317    domain: Domains,
318    range: Option<Domains>,
319}
320
321impl Signature for Type23Signature {
322    fn clip_returns(&self, returns: FunctionValue) -> FunctionValue {
323        let Some(range) = self.range.as_ref() else {
324            return returns;
325        };
326        debug_assert_eq!(returns.len(), range.n());
327
328        returns
329            .iter()
330            .zip(range.0.iter())
331            .map(|(&ret, domain)| domain.clamp(ret))
332            .collect()
333    }
334
335    fn clip_args(&self, args: &[f32]) -> TinyVec<[f32; 4]> {
336        debug_assert_eq!(args.len(), self.n_args());
337
338        args.iter()
339            .zip(self.domain.0.iter())
340            .map(|(&arg, domain)| domain.clamp(arg))
341            .collect()
342    }
343}
344
345impl Type23Signature {
346    pub fn new(domain: Domains, range: Option<Domains>) -> Self {
347        Self { domain, range }
348    }
349
350    pub fn n_args(&self) -> usize {
351        self.domain.n()
352    }
353
354    pub fn n_returns(&self) -> Option<usize> {
355        self.range.as_ref().map(Domains::n)
356    }
357}
358
359/// Function signature for Type 0 and 4, which range is required
360#[derive(Debug, PartialEq, Clone)]
361pub struct Type04Signature {
362    domain: Domains,
363    range: Domains,
364}
365
366impl Signature for Type04Signature {
367    fn clip_returns(&self, returns: FunctionValue) -> FunctionValue {
368        debug_assert_eq!(returns.len(), self.range.n());
369
370        returns
371            .iter()
372            .zip(self.range.0.iter())
373            .map(|(&ret, domain)| domain.clamp(ret))
374            .collect()
375    }
376
377    fn clip_args(&self, args: &[f32]) -> TinyVec<[f32; 4]> {
378        debug_assert_eq!(args.len(), self.n_args());
379
380        args.iter()
381            .zip(self.domain.0.iter())
382            .map(|(&arg, domain)| domain.clamp(arg))
383            .collect()
384    }
385}
386
387impl Type04Signature {
388    pub fn new(domain: Domains, range: Domains) -> Self {
389        Self { domain, range }
390    }
391
392    pub fn n_args(&self) -> usize {
393        self.domain.n()
394    }
395
396    pub fn n_returns(&self) -> usize {
397        self.range.n()
398    }
399}
400
401fn f32_zero_arr() -> Vec<f32> {
402    vec![0.0]
403}
404
405fn f32_one_arr() -> Vec<f32> {
406    vec![1.0]
407}
408
409#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromIntObject, Default)]
410pub enum InterpolationOrder {
411    #[default]
412    Linear = 1,
413    Cubic = 3,
414}
415
416#[pdf_object(0i32)]
417#[type_field("FunctionType")]
418#[root_pdf_object]
419pub trait SampledFunctionDictTrait {
420    fn size(&self) -> Vec<u32>;
421    fn bits_per_sample(&self) -> u32;
422
423    #[try_from]
424    #[or_default]
425    fn order(&self) -> InterpolationOrder;
426
427    #[try_from]
428    fn encode(&self) -> Option<Domains>;
429
430    #[try_from]
431    fn decode(&self) -> Option<Domains>;
432
433    #[try_from]
434    fn domain(&self) -> Domains;
435
436    #[try_from]
437    fn range(&self) -> Option<Domains>;
438}
439
440/// struct to implement Function trait for SampledFunctionDict,
441/// because sampled function need to load sample data from stream.
442#[derive(Debug, PartialEq, Clone)]
443pub struct SampledFunction {
444    signature: Type04Signature,
445    encode: Domains,
446    decode: Domains,
447    size: Vec<u32>,
448    samples: Vec<u8>,
449    bits_per_sample: u8,
450}
451
452impl SampledFunction {
453    fn samples(&self) -> usize {
454        // NOTE: assume bits_per_sample is 8
455        self.samples.len() / self.signature.n_returns()
456    }
457}
458
459impl InnerFunction for SampledFunction {
460    type Signature = Type04Signature;
461
462    fn inner_call(&self, args: TinyVec<[f32; 4]>) -> Result<FunctionValue> {
463        let mut idx = 0;
464        for (arg, (domain, (encode, size))) in args
465            .iter()
466            .zip(
467                self.signature
468                    .domain
469                    .iter()
470                    .zip(self.encode.iter().zip(self.size.iter())),
471            )
472            .rev()
473        {
474            let arg = (arg - domain.start) / (domain.end - domain.start);
475            let arg = arg.mul_add(encode.end - encode.start, encode.start);
476            idx = size * idx
477                + arg
478                    .round()
479                    .to_u32()
480                    .whatever_context::<_, ObjectValueError>("convert to u32")?
481                    .clamp(0, *size - 1);
482        }
483        let idx = idx as usize;
484
485        let n_ret = self.signature.n_returns();
486        let sample_size = self.bits_per_sample as usize / 8;
487        let mut r = tiny_vec![];
488        let decode = &self.decode.0[0];
489        for i in 0..n_ret {
490            let start_p = (idx * n_ret + i) * sample_size;
491            let mut sample = 0u32;
492            for i in 0..sample_size {
493                sample <<= 8;
494                sample |= self.samples[start_p + i] as u32;
495            }
496            let sample = sample as f32 / (2.0_f32.powi(self.bits_per_sample as i32) - 1.0);
497            let sample = sample.mul_add(decode.end - decode.start, decode.start);
498            r.push(sample);
499        }
500        Ok(r)
501    }
502
503    fn signature(&self) -> &Self::Signature {
504        &self.signature
505    }
506
507    fn stops(&self) -> impl Iterator<Item = f32> + 'static {
508        let domain = &self.signature.domain.0[0];
509        let samples = self.samples().min(256);
510        let t0 = euclid::default::Length::new(domain.start);
511        let t1 = euclid::default::Length::new(domain.end);
512
513        (0..samples).map(move |i| t0.lerp(t1, i as f32 / (samples - 1) as f32).0)
514    }
515}
516
517impl SampledFunctionDict<'_, '_> {
518    fn type04_signature(&self) -> Result<Type04Signature> {
519        Ok(Type04Signature {
520            domain: self.domain()?,
521            range: self
522                .range()?
523                .whatever_context::<_, ObjectValueError>("range should exist")?,
524        })
525    }
526
527    /// Return SampledFunction instance which implements Function trait.
528    pub fn func(&self) -> Result<SampledFunction> {
529        let bits_per_sample = self.bits_per_sample()?;
530        ensure_whatever!(bits_per_sample >= 8, "todo: support bits_per_sample < 8");
531        ensure_whatever!(
532            InterpolationOrder::Linear == self.order()?,
533            "todo: support cubic interpolation"
534        );
535
536        let size = self.size()?;
537        let resolver = self.d.resolver();
538        let stream = resolver
539            .resolve(self.id)
540            .whatever_context::<_, ObjectValueError>("resolve object")?
541            .as_stream()
542            .whatever_context::<_, ObjectValueError>("get as stream")?;
543        let sample_data = stream
544            .decode(resolver)
545            .whatever_context::<_, ObjectValueError>("decode stream")?;
546        let signature = self.type04_signature()?;
547        ensure_whatever!(
548            sample_data.len() >= size[0] as usize * signature.n_returns(),
549            "Sample data length is insufficient"
550        );
551        Ok(SampledFunction {
552            signature,
553            encode: self.encode()?.unwrap_or_else(|| {
554                Domains(
555                    size.iter()
556                        .map(|v| Domain::new(0.0, (*v - 1) as f32))
557                        .collect(),
558                )
559            }),
560            decode: self.decode()?.map_or_else(
561                || {
562                    self.range()
563                        .whatever_context::<_, ObjectValueError>("get range")?
564                        .whatever_context::<_, ObjectValueError>(
565                            "range should exist in sampled function",
566                        )
567                },
568                Ok,
569            )?,
570            size: self.size()?,
571            samples: sample_data.into_owned(),
572            bits_per_sample: bits_per_sample.try_into().unwrap_or(u8::MAX),
573        })
574    }
575}
576
577#[pdf_object(2i32)]
578#[type_field("FunctionType")]
579pub trait ExponentialInterpolationFunctionDictTrait {
580    #[default_fn(f32_zero_arr)]
581    fn c0(&self) -> Vec<f32>;
582
583    #[default_fn(f32_one_arr)]
584    fn c1(&self) -> Vec<f32>;
585
586    fn n(&self) -> f32;
587
588    #[try_from]
589    fn domain(&self) -> Domains;
590
591    #[try_from]
592    fn range(&self) -> Option<Domains>;
593}
594
595pub struct ExponentialInterpolationFunction {
596    c0: Vec<f32>,
597    c1: Vec<f32>,
598    n: f32,
599    signature: Type23Signature,
600}
601
602impl InnerFunction for ExponentialInterpolationFunction {
603    type Signature = Type23Signature;
604
605    fn inner_call(&self, args: TinyVec<[f32; 4]>) -> Result<FunctionValue> {
606        let x = args[0];
607        let r = (0..self.c0.len())
608            .map(|i| x.powf(self.n).mul_add(self.c1[i] - self.c0[i], self.c0[i]))
609            .collect();
610        Ok(r)
611    }
612
613    fn signature(&self) -> &Type23Signature {
614        &self.signature
615    }
616
617    fn stops(&self) -> impl Iterator<Item = f32> + 'static {
618        // For exponential interpolation, we just need start and end points
619        // of the domain, similar to how build_stops() handles it
620        let domain = &self.signature.domain.0[0]; // Get first domain
621        vec![domain.start, domain.end].into_iter()
622    }
623}
624
625impl ExponentialInterpolationFunctionDict<'_, '_> {
626    fn func(&self) -> Result<ExponentialInterpolationFunction> {
627        Ok(ExponentialInterpolationFunction {
628            c0: self.c0()?,
629            c1: self.c1()?,
630            n: self.n()?,
631            signature: self.type23_signature()?,
632        })
633    }
634
635    fn type23_signature(&self) -> Result<Type23Signature> {
636        Ok(Type23Signature {
637            domain: self.domain()?,
638            range: self.range()?,
639        })
640    }
641}
642
643#[pdf_object(3i32)]
644#[type_field("FunctionType")]
645pub trait StitchingFunctionDictTrait {
646    /// The number of values shall be `k - 1`
647    fn bounds(&self) -> Vec<f32>;
648
649    /// The number of values shall be `k`
650    #[try_from]
651    fn encode(&self) -> Domains;
652
653    #[try_from]
654    fn domain(&self) -> Domains;
655
656    #[try_from]
657    fn range(&self) -> Option<Domains>;
658}
659
660impl StitchingFunctionDict<'_, '_> {
661    fn func(&self) -> Result<StitchingFunction> {
662        let functions: Vec<Box<dyn Function>> =
663            self.d
664                .zero_one_or_more(&sname("Functions"))
665                .whatever_context::<_, ObjectValueError>("get stitching Functions")?;
666        let bounds = self.bounds()?;
667        let encode = self.encode()?;
668        let signature = Type23Signature {
669            domain: self.domain()?,
670            range: self.range()?,
671        };
672        Ok(StitchingFunction {
673            functions,
674            bounds,
675            encode,
676            signature,
677        })
678    }
679}
680
681pub struct StitchingFunction {
682    functions: Vec<Box<dyn Function>>,
683    bounds: Vec<f32>,
684    encode: Domains,
685    signature: Type23Signature,
686}
687
688impl StitchingFunction {
689    fn find_function(bounds: &[f32], x: f32) -> usize {
690        bounds
691            .iter()
692            .position(|&bound| x < bound)
693            .unwrap_or(bounds.len())
694    }
695
696    fn sub_domain(domain: Domain, bounds: &[f32], idx: usize) -> Domain {
697        let start = if idx == 0 {
698            domain.start
699        } else {
700            bounds[idx - 1]
701        };
702        let end = if idx == bounds.len() {
703            domain.end
704        } else {
705            bounds[idx]
706        };
707        Domain::new(start, end)
708    }
709
710    fn interpolation(from: Domain, to: Domain, t: f32) -> f32 {
711        let a_len = from.end - from.start;
712        let b_len = to.end - to.start;
713        let t = (t - from.start) / a_len;
714        t.mul_add(b_len, to.start) // t * b_len + b.start
715    }
716
717    fn domains(&self) -> &Domains {
718        &self.signature.domain
719    }
720}
721
722impl InnerFunction for StitchingFunction {
723    type Signature = Type23Signature;
724
725    fn inner_call(&self, args: TinyVec<[f32; 4]>) -> Result<FunctionValue> {
726        ensure_whatever!(args.len() == 1, "expected one argument");
727
728        let x = args[0];
729        let function_idx = Self::find_function(&self.bounds, x);
730        let mut sub_domain = Self::sub_domain(self.domains().0[0], &self.bounds, function_idx);
731        if sub_domain.is_zero() {
732            // possibly incorrect bounds, bounds[0] should > domain[0].start
733            // bounds[last] should < domain[0].end, but some buggie file
734            // breaks, cause a zero sub_domain
735            sub_domain = self.domains().0[0];
736        }
737        let x1 = Self::interpolation(sub_domain, self.encode.0[function_idx], x);
738
739        let f = &self.functions[function_idx];
740        let r = f.call(&[x1])?;
741        Ok(r)
742    }
743
744    fn signature(&self) -> &Type23Signature {
745        &self.signature
746    }
747
748    fn stops(&self) -> impl Iterator<Item = f32> + 'static {
749        let domain = self.domains().0[0];
750        let mut stops = Vec::with_capacity(self.bounds.len() + 2);
751        stops.push(domain.start);
752        for t in &self.bounds {
753            stops.push(*t);
754        }
755        stops.push(domain.end);
756        stops.into_iter()
757    }
758}
759
760#[cfg(test)]
761mod tests;