1use crate::array::Array;
2use crate::*;
3use core::ops::{Index, IndexMut};
4
5pub struct ArrayMap<K: ArrayFinite<V>, V>(K::Array);
34
35#[doc(hidden)]
42#[allow(clippy::missing_safety_doc)] pub unsafe trait ArrayFinite<V>: Finite {
44 #[allow(missing_docs)]
45 type Array: Array<V>;
46}
47
48impl<K: ArrayFinite<V>, V> ArrayMap<K, V> {
49 pub fn new(mut f: impl FnMut(K) -> V) -> Self {
51 ArrayMap(K::Array::new(|k| {
52 f(unsafe { K::nth(k).unwrap_unchecked() })
53 }))
54 }
55
56 pub fn from(array: K::Array) -> Self {
67 Self(array)
68 }
69
70 pub fn map_with_key<N>(&self, mut f: impl FnMut(K, &V) -> N) -> ArrayMap<K, N>
72 where
73 K: ArrayFinite<N>,
74 {
75 ArrayMap(<K as ArrayFinite<N>>::Array::new(|k| unsafe {
76 f(
77 K::nth(k).unwrap_unchecked(),
78 self.0.as_slice().get_unchecked(k),
79 )
80 }))
81 }
82
83 pub fn map<N>(&self, mut f: impl FnMut(&V) -> N) -> ArrayMap<K, N>
85 where
86 K: ArrayFinite<N>,
87 {
88 self.map_with_key(|_, v| f(v))
89 }
90}
91
92impl<K: ArrayFinite<V>, V: Default> Default for ArrayMap<K, V> {
93 fn default() -> Self {
94 ArrayMap(K::Array::new(|_| Default::default()))
95 }
96}
97
98impl<K: ArrayFinite<V>, V> Index<K> for ArrayMap<K, V> {
99 type Output = V;
100 fn index(&self, index: K) -> &Self::Output {
101 let index = K::index_of(index);
102 unsafe { self.0.as_slice().get_unchecked(index) }
103 }
104}
105
106impl<K: ArrayFinite<V>, V> IndexMut<K> for ArrayMap<K, V> {
107 fn index_mut(&mut self, index: K) -> &mut Self::Output {
108 let index = K::index_of(index);
109 unsafe { self.0.as_slice_mut().get_unchecked_mut(index) }
110 }
111}
112
113impl<K: CompressFinite + ArrayFinite<V>, V> Index<Compress<K>> for ArrayMap<K, V> {
114 type Output = V;
115 fn index(&self, index: Compress<K>) -> &Self::Output {
116 let index = Compress::index_of(index);
117 unsafe { self.0.as_slice().get_unchecked(index) }
118 }
119}
120
121impl<K: CompressFinite + ArrayFinite<V>, V> IndexMut<Compress<K>> for ArrayMap<K, V> {
122 fn index_mut(&mut self, index: Compress<K>) -> &mut Self::Output {
123 let index = Compress::index_of(index);
124 unsafe { self.0.as_slice_mut().get_unchecked_mut(index) }
125 }
126}
127
128impl<K: ArrayFinite<V>, V> Clone for ArrayMap<K, V>
129where
130 K::Array: Clone,
131{
132 fn clone(&self) -> Self {
133 Self(self.0.clone())
134 }
135}
136
137impl<K: ArrayFinite<V>, V> Copy for ArrayMap<K, V> where K::Array: Copy {}
138
139impl<K: ArrayFinite<V>, V> PartialEq for ArrayMap<K, V>
140where
141 K::Array: PartialEq,
142{
143 fn eq(&self, other: &Self) -> bool {
144 self.0 == other.0
145 }
146}
147
148impl<K: ArrayFinite<V>, V> Eq for ArrayMap<K, V> where K::Array: Eq {}
149
150impl<K: ArrayFinite<V>, V> PartialOrd for ArrayMap<K, V>
151where
152 K::Array: PartialOrd,
153{
154 fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
155 self.0.partial_cmp(&other.0)
156 }
157}
158
159impl<K: ArrayFinite<V>, V> Ord for ArrayMap<K, V>
160where
161 K::Array: Ord,
162{
163 fn cmp(&self, other: &Self) -> core::cmp::Ordering {
164 self.0.cmp(&other.0)
165 }
166}
167
168#[test]
169fn test_map_with_key() {
170 let map = ArrayMap::new(|x| if x { 1 } else { 0 });
171 let map = map.map_with_key(|k, v| if k { *v * 2 } else { *v + 5 });
172 assert_eq!(map[false], 5);
173 assert_eq!(map[true], 2);
174}