1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use crate::internal::*;
use std::fmt;

mod lir;
mod mir;

pub use lir::{LirScan, State};
pub use mir::Scan;

#[derive(Clone, new, Hash, Eq, PartialEq, Copy, Debug)]
pub struct ScanInfo {
    pub slot: usize,
    pub axis: usize,
    pub chunk: isize,
}

#[derive(Clone, new, Hash)]
pub enum InputMapping {
    Full { slot: usize },
    State { initializer: StateInitializer },
    Scan(ScanInfo),
}

impl InputMapping {
    pub fn as_state(&self) -> Option<&StateInitializer> {
        match self {
            InputMapping::State { initializer } => Some(initializer),
            _ => None,
        }
    }

    pub fn as_scan(&self) -> Option<&ScanInfo> {
        match self {
            InputMapping::Scan(s) => Some(s),
            _ => None,
        }
    }

    pub fn invisible(&self) -> bool {
        matches!(self, InputMapping::State { initializer: StateInitializer::Value(_) })
    }

    pub fn slot(&self) -> Option<usize> {
        match self {
            InputMapping::Full { slot } => Some(*slot),
            InputMapping::Scan(info) => Some(info.slot),
            InputMapping::State { initializer } => match initializer {
                StateInitializer::FromInput(slot) => Some(*slot),
                _ => None,
            },
        }
    }
}

impl fmt::Debug for InputMapping {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match self {
            InputMapping::Full { slot } => write!(fmt, "Full, inlet {slot}"),
            InputMapping::State { initializer } => {
                write!(fmt, "State initialized by {initializer:?}")
            }
            InputMapping::Scan(info) => {
                write!(
                    fmt,
                    "Scan inlet {}, axis: {}, chunk: {:?}.",
                    info.slot, info.axis, info.chunk
                )
            }
        }
    }
}

#[derive(Clone, new, Hash, Default)]
pub struct OutputMapping<F: Clone> {
    pub scan: Option<ScanInfo>,
    pub full_dim_hint: Option<F>,
    pub last_value_slot: Option<usize>,
    pub state: bool,
}

impl<F: Clone> OutputMapping<F> {
    pub fn invisible(&self) -> bool {
        self.scan.is_none() && self.last_value_slot.is_none()
    }
}

impl<F: Clone + DimLike> OutputMapping<F> {
    pub fn concretize_dims(&self, values: &SymbolValues) -> TractResult<OutputMapping<F>> {
        Ok(Self {
            full_dim_hint: self.full_dim_hint.as_ref().map(|h| h.eval(values)),
            ..self.clone()
        })
    }
}

impl<F: Clone + fmt::Display> fmt::Debug for OutputMapping<F> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.state {
            write!(fmt, "State. ")?;
        }
        if let Some(last_value_slot) = self.last_value_slot {
            write!(fmt, "Last value to outlet {last_value_slot}. ")?;
        }
        if let Some(info) = self.scan {
            write!(fmt, "Full value to outlet {} (axis: {}). ", info.slot, info.axis)?;
        }
        if let Some(full_dim_hint) = &self.full_dim_hint {
            write!(fmt, "Full len {full_dim_hint}. ")?;
        }
        Ok(())
    }
}

#[derive(Clone, new, Hash)]
pub enum StateInitializer {
    FromInput(usize),
    Value(Arc<Tensor>),
}

impl fmt::Debug for StateInitializer {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        use StateInitializer::*;
        match self {
            FromInput(i) => write!(fmt, "inlet {i}"),
            Value(t) => write!(fmt, "tensor {t:?}"),
        }
    }
}