Skip to main content

tenferro_runtime/program/
bindings.rs

1use std::fmt;
2use std::sync::Arc;
3
4use crate::checkpoint::RetainedValue;
5
6use super::value::ProgramBuilderNonce;
7use super::{BindingKey, ProgramValue};
8
9#[derive(Clone)]
10pub(crate) struct PendingBinding {
11    pub(crate) key: BindingKey,
12    pub(crate) input: ProgramValue,
13    pub(crate) tensor: Arc<RetainedValue>,
14}
15
16struct ProgramBinding {
17    key: BindingKey,
18    tensor: Arc<RetainedValue>,
19}
20
21/// Immutable tensor defaults and large constants kept outside semantic structure.
22#[derive(Clone)]
23pub struct ProgramBindings {
24    owner: ProgramBuilderNonce,
25    entries: Arc<[ProgramBinding]>,
26}
27
28impl ProgramBindings {
29    pub(crate) fn freeze(owner: ProgramBuilderNonce, pending: Vec<PendingBinding>) -> Self {
30        let mut entries: Vec<_> = pending
31            .into_iter()
32            .map(|binding| ProgramBinding {
33                key: binding.key,
34                tensor: binding.tensor,
35            })
36            .collect();
37        entries.sort_unstable_by_key(|entry| entry.key.slot);
38        Self {
39            owner,
40            entries: entries.into(),
41        }
42    }
43
44    pub(crate) fn belongs_to(&self, owner: ProgramBuilderNonce) -> bool {
45        self.owner == owner
46    }
47
48    pub(crate) fn tensor_for_input(&self, input: ProgramValue) -> Option<Arc<RetainedValue>> {
49        self.tensor_ref_for_input(input).map(Arc::clone)
50    }
51
52    pub(crate) fn tensor_ref_for_input(&self, input: ProgramValue) -> Option<&Arc<RetainedValue>> {
53        if input.owner != self.owner {
54            return None;
55        }
56        self.entries
57            .binary_search_by_key(&input.slot, |entry| entry.key.slot)
58            .ok()
59            .and_then(|index| self.entries.get(index))
60            .map(|entry| &entry.tensor)
61    }
62
63    pub(crate) fn bound_inputs(&self) -> impl ExactSizeIterator<Item = ProgramValue> + '_ {
64        self.entries
65            .iter()
66            .map(|entry| ProgramValue::new(entry.key.slot, self.owner))
67    }
68
69    #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
70    pub(crate) fn preserves_all(&self, source: &Self) -> bool {
71        if self.entries.len() < source.entries.len() {
72            return false;
73        }
74        let mut matched = vec![false; self.entries.len()];
75        source.entries.iter().all(|source_entry| {
76            self.entries
77                .iter()
78                .enumerate()
79                .find(|(index, entry)| {
80                    !matched[*index] && Arc::ptr_eq(&entry.tensor, &source_entry.tensor)
81                })
82                .map(|(index, _)| {
83                    matched[index] = true;
84                })
85                .is_some()
86        })
87    }
88
89    #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
90    pub(crate) fn cache_exact_eq(&self, other: &Self) -> bool {
91        self.entries.len() == other.entries.len()
92            && self
93                .entries
94                .iter()
95                .zip(other.entries.iter())
96                .all(|(left, right)| {
97                    left.key.slot == right.key.slot && Arc::ptr_eq(&left.tensor, &right.tensor)
98                })
99    }
100
101    /// Return the number of tensor bindings.
102    pub fn len(&self) -> usize {
103        self.entries.len()
104    }
105
106    /// Return whether no tensor bindings are present.
107    pub fn is_empty(&self) -> bool {
108        self.entries.is_empty()
109    }
110
111    /// Borrow a tensor by opaque binding key.
112    pub fn get(&self, key: BindingKey) -> Option<&RetainedValue> {
113        if key.owner != self.owner {
114            return None;
115        }
116        self.entries
117            .binary_search_by_key(&key.slot, |entry| entry.key.slot)
118            .ok()
119            .and_then(|index| self.entries.get(index))
120            .filter(|entry| entry.key == key)
121            .map(|entry| entry.tensor.as_ref())
122    }
123
124    /// Iterate over ordered binding keys and borrowed move-only values.
125    pub fn iter(&self) -> impl ExactSizeIterator<Item = (BindingKey, &RetainedValue)> + '_ {
126        self.entries
127            .iter()
128            .map(|entry| (entry.key, entry.tensor.as_ref()))
129    }
130}
131
132impl fmt::Debug for ProgramBindings {
133    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
134        formatter
135            .debug_struct("ProgramBindings")
136            .field("len", &self.entries.len())
137            .finish()
138    }
139}