as_derive_utils/datastructure/
field_map.rs

1use super::*;
2
3use std::{
4    mem,
5    ops::{Index, IndexMut},
6};
7
8/// This is a map from fields to some value.
9///
10/// If you put this in a type,and use Default to initialize it,
11/// you must remember to replace the `FieldMap` using either `FieldMap::defaulted` or
12/// `FieldMap::with`
13#[derive(Default, Clone, Debug, PartialEq, Hash)]
14pub struct FieldMap<T> {
15    // The outer vec is the enum variant (if it's a struct/union it's a single element Vec),
16    // the inner one is the field within a variant/struct/union.
17    fields: Vec<Vec<T>>,
18}
19
20impl<T> FieldMap<T> {
21    /// Constructs an empty FieldMap.
22    pub fn empty() -> Self {
23        Self { fields: Vec::new() }
24    }
25
26    /// Constructs an FieldMap which maps each field in the DataStructure
27    /// to the default value for `T`.
28    pub fn defaulted<'a>(ds: &'a DataStructure<'a>) -> Self
29    where
30        T: Default,
31    {
32        Self::with(ds, |_| T::default())
33    }
34
35    /// Constructs an FieldMap which maps each field in the DataStructure to a value
36    /// (obtained by mapping each individual field to `T` using a closure).
37    pub fn with<'a, F>(ds: &'a DataStructure<'a>, mut f: F) -> Self
38    where
39        F: FnMut(&'a Field<'a>) -> T,
40    {
41        Self {
42            fields: ds
43                .variants
44                .iter()
45                .map(|vari| vari.fields.iter().map(&mut f).collect::<Vec<_>>())
46                .collect::<Vec<_>>(),
47        }
48    }
49
50    /// Maps each value in the map to another one,using a closure.
51    pub fn map<F, U>(self, mut f: F) -> FieldMap<U>
52    where
53        F: FnMut(FieldIndex, T) -> U,
54    {
55        let fields = self
56            .fields
57            .into_iter()
58            .enumerate()
59            .map(|(var_i, variant)| {
60                variant
61                    .into_iter()
62                    .enumerate()
63                    .map(|(pos, v)| {
64                        let index = FieldIndex {
65                            variant: var_i,
66                            pos,
67                        };
68                        f(index, v)
69                    })
70                    .collect::<Vec<U>>()
71            })
72            .collect::<Vec<Vec<U>>>();
73        FieldMap { fields }
74    }
75
76    /// Whether the field index maps to a field.
77    #[allow(dead_code)]
78    pub fn contains_index(&self, index: FieldIndex) -> bool {
79        self.fields
80            .get(index.variant)
81            .map_or(false, |variant| index.pos < variant.len())
82    }
83
84    /// Add a new field to the map along with a value that it maps into.
85    pub fn insert(&mut self, field: &Field<'_>, value: T) -> T {
86        mem::replace(&mut self[field], value)
87    }
88
89    pub fn iter(&self) -> impl Iterator<Item = (FieldIndex, &'_ T)> + Clone + '_ {
90        self.fields.iter().enumerate().flat_map(|(v_i, v)| {
91            v.iter().enumerate().map(move |(f_i, f)| {
92                let index = FieldIndex {
93                    variant: v_i as _,
94                    pos: f_i as _,
95                };
96                (index, f)
97            })
98        })
99    }
100
101    pub fn iter_mut(&mut self) -> impl Iterator<Item = (FieldIndex, &'_ mut T)> + '_ {
102        self.fields.iter_mut().enumerate().flat_map(|(v_i, v)| {
103            v.iter_mut().enumerate().map(move |(f_i, f)| {
104                let index = FieldIndex {
105                    variant: v_i as _,
106                    pos: f_i as _,
107                };
108                (index, f)
109            })
110        })
111    }
112
113    pub fn values(&self) -> impl Iterator<Item = &'_ T> + Clone + '_ {
114        self.fields.iter().flat_map(|v| v.iter())
115    }
116}
117
118impl<T> Index<FieldIndex> for FieldMap<T> {
119    type Output = T;
120
121    fn index(&self, index: FieldIndex) -> &T {
122        &self.fields[index.variant][index.pos]
123    }
124}
125
126impl<T> IndexMut<FieldIndex> for FieldMap<T> {
127    fn index_mut(&mut self, index: FieldIndex) -> &mut T {
128        &mut self.fields[index.variant][index.pos]
129    }
130}
131
132impl<'a, 'b, T> Index<&'a Field<'b>> for FieldMap<T> {
133    type Output = T;
134
135    fn index(&self, field: &'a Field<'b>) -> &T {
136        let index = field.index;
137        &self.fields[index.variant][index.pos]
138    }
139}
140
141impl<'a, 'b, T> IndexMut<&'a Field<'b>> for FieldMap<T> {
142    fn index_mut(&mut self, field: &'a Field<'b>) -> &mut T {
143        let index = field.index;
144        &mut self.fields[index.variant][index.pos]
145    }
146}