Skip to main content

wxy_domain/
domain.rs

1//! Domains: index space plus the mappings that place it.
2
3use crate::dim::Dim;
4use crate::mapping::Mapping;
5use crate::positions::Rung;
6
7/// Where elements are: index space, and any number of mappings.
8///
9/// A domain keeps every mapping it was given (design doc 02, gate 2
10/// default 2). One may be marked primary per dim; [`Domain::primary_for`]
11/// returns it, falling back to the first mapping that covers the dim.
12#[derive(Debug, Clone, Default)]
13pub struct Domain {
14    pub dims: Vec<Dim>,
15    pub mappings: Vec<Mapping>,
16    /// Indices into `mappings` marked primary. At most one per dim.
17    pub primary: Vec<usize>,
18}
19
20impl Domain {
21    pub fn new(dims: Vec<Dim>) -> Self {
22        Self {
23            dims,
24            ..Default::default()
25        }
26    }
27
28    pub fn shape(&self) -> Vec<u64> {
29        self.dims.iter().map(|d| d.size).collect()
30    }
31
32    pub fn dim_index(&self, name: &str) -> Option<usize> {
33        self.dims.iter().position(|d| d.name == name)
34    }
35
36    /// Add a mapping and return its index. Panics if it names a dim the
37    /// domain does not have.
38    pub fn add_mapping(&mut self, mapping: Mapping) -> usize {
39        assert!(
40            mapping.dims.iter().all(|&d| d < self.dims.len()),
41            "mapping covers a dim outside the domain"
42        );
43        self.mappings.push(mapping);
44        self.mappings.len() - 1
45    }
46
47    /// Mark a mapping primary for the dims it covers, replacing any other
48    /// primary mapping on those dims.
49    pub fn set_primary(&mut self, index: usize) {
50        let dims = self.mappings[index].dims.clone();
51        let mappings = &self.mappings;
52        self.primary
53            .retain(|&p| !mappings[p].dims.iter().any(|d| dims.contains(d)));
54        self.primary.push(index);
55    }
56
57    /// Mappings that cover dim `d`.
58    pub fn mappings_for(&self, d: usize) -> impl Iterator<Item = (usize, &Mapping)> {
59        self.mappings
60            .iter()
61            .enumerate()
62            .filter(move |(_, m)| m.dims.contains(&d))
63    }
64
65    /// The primary mapping for dim `d`.
66    pub fn primary_for(&self, d: usize) -> Option<&Mapping> {
67        self.primary
68            .iter()
69            .map(|&p| &self.mappings[p])
70            .find(|m| m.dims.contains(&d))
71            .or_else(|| self.mappings_for(d).map(|(_, m)| m).next())
72    }
73
74    /// The rung of the primary mapping for dim `d` (R0 when unmapped).
75    pub fn rung_for(&self, d: usize) -> Rung {
76        self.primary_for(d).map_or(Rung::IndexOnly, Mapping::rung)
77    }
78}