1use crate::dim::Dim;
4use crate::mapping::Mapping;
5use crate::positions::Rung;
6
7#[derive(Debug, Clone, Default)]
13pub struct Domain {
14 pub dims: Vec<Dim>,
15 pub mappings: Vec<Mapping>,
16 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 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 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 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 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 pub fn rung_for(&self, d: usize) -> Rung {
76 self.primary_for(d).map_or(Rung::IndexOnly, Mapping::rung)
77 }
78}