pintos-os.org Git - pspp/blob - rust/src/dictionary.rs

   1 use std::{
   2     collections::{HashMap, HashSet},
   3     fmt::Debug,
   4 };
   5
   6 use encoding_rs::Encoding;
   7 use indexmap::IndexSet;
   8
   9 use crate::{
  10     cooked::{Alignment, Measure, MissingValues, Value, VarWidth},
  11     format::Format,
  12     identifier::{ByIdentifier, HasIdentifier, Identifier},
  13     raw::CategoryLabels,
  14 };
  15
  16 pub type DictIndex = usize;
  17
  18 #[derive(Clone, Debug)]
  19 pub struct Dictionary {
  20     pub variables: IndexSet<ByIdentifier<Variable>>,
  21     pub split_file: Vec<DictIndex>,
  22     pub weight: Option<DictIndex>,
  23     pub filter: Option<DictIndex>,
  24     pub case_limit: Option<u64>,
  25     pub file_label: Option<String>,
  26     pub documents: Vec<String>,
  27     pub vectors: HashSet<ByIdentifier<Vector>>,
  28     pub attributes: HashSet<ByIdentifier<Attribute>>,
  29     pub mrsets: HashSet<ByIdentifier<MultipleResponseSet>>,
  30     pub variable_sets: HashSet<ByIdentifier<VariableSet>>,
  31     pub encoding: &'static Encoding,
  32 }
  33
  34 impl Dictionary {
  35     pub fn new(encoding: &'static Encoding) -> Self {
  36         Self {
  37             variables: IndexSet::new(),
  38             split_file: Vec::new(),
  39             weight: None,
  40             filter: None,
  41             case_limit: None,
  42             file_label: None,
  43             documents: Vec::new(),
  44             vectors: HashSet::new(),
  45             attributes: HashSet::new(),
  46             mrsets: HashSet::new(),
  47             variable_sets: HashSet::new(),
  48             encoding,
  49         }
  50     }
  51
  52     pub fn delete_vars(&mut self, start: DictIndex, count: usize) {
  53         self.update_dict_indexes(&|index| {
  54             if index < start {
  55                 Some(index)
  56             } else if index < start + count {
  57                 None
  58             } else {
  59                 Some(index - count)
  60             }
  61         })
  62     }
  63
  64     fn update_dict_indexes<F>(&mut self, f: &F)
  65     where
  66         F: Fn(DictIndex) -> Option<DictIndex>,
  67     {
  68         update_dict_index_vec(&mut self.split_file, f);
  69         self.weight = self.weight.map(|index| f(index)).flatten();
  70         self.filter = self.filter.map(|index| f(index)).flatten();
  71         self.vectors = self
  72             .vectors
  73             .drain()
  74             .filter_map(|vector_by_id| {
  75                 vector_by_id
  76                     .0
  77                     .with_updated_dict_indexes(f)
  78                     .map(|vector| ByIdentifier::new(vector))
  79             })
  80             .collect();
  81         self.mrsets = self
  82             .mrsets
  83             .drain()
  84             .filter_map(|mrset_by_id| {
  85                 mrset_by_id
  86                     .0
  87                     .with_updated_dict_indexes(f)
  88                     .map(|mrset| ByIdentifier::new(mrset))
  89             })
  90             .collect();
  91         self.variable_sets = self
  92             .variable_sets
  93             .drain()
  94             .filter_map(|var_set_by_id| {
  95                 var_set_by_id
  96                     .0
  97                     .with_updated_dict_indexes(f)
  98                     .map(|var_set| ByIdentifier::new(var_set))
  99             })
 100             .collect();
 101     }
 102 }
 103
 104 fn update_dict_index_vec<F>(dict_indexes: &mut Vec<DictIndex>, f: F)
 105 where
 106     F: Fn(DictIndex) -> Option<DictIndex>,
 107 {
 108     dict_indexes.retain_mut(|index| {
 109         if let Some(new) = f(*index) {
 110             *index = new;
 111             true
 112         } else {
 113             false
 114         }
 115     });
 116 }
 117
 118 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 119 pub enum Role {
 120     Input,
 121     Target,
 122     Both,
 123     None,
 124     Partition,
 125     Split,
 126 }
 127
 128 #[derive(Clone, Debug)]
 129 pub struct Variable {
 130     pub name: Identifier,
 131     pub width: VarWidth,
 132     pub missing_values: MissingValues,
 133     pub print_format: Format,
 134     pub write_format: Format,
 135     pub value_labels: HashMap<Value, String>,
 136     pub label: Option<String>,
 137     pub measure: Measure,
 138     pub role: Role,
 139     pub display_width: u32,
 140     pub alignment: Alignment,
 141     pub leave: bool,
 142     pub short_names: Vec<Identifier>,
 143     pub attributes: HashSet<ByIdentifier<Attribute>>,
 144 }
 145
 146 impl HasIdentifier for Variable {
 147     fn identifier(&self) -> &Identifier {
 148         &self.name
 149     }
 150 }
 151
 152 #[derive(Clone, Debug)]
 153 pub struct Vector {
 154     pub name: Identifier,
 155     pub variables: Vec<DictIndex>,
 156 }
 157
 158 impl Vector {
 159     fn with_updated_dict_indexes(
 160         mut self,
 161         f: impl Fn(DictIndex) -> Option<DictIndex>,
 162     ) -> Option<Self> {
 163         update_dict_index_vec(&mut self.variables, f);
 164         (!self.variables.is_empty()).then_some(self)
 165     }
 166 }
 167
 168 impl HasIdentifier for Vector {
 169     fn identifier(&self) -> &Identifier {
 170         &self.name
 171     }
 172 }
 173
 174 #[derive(Clone, Debug)]
 175 pub struct Attribute {
 176     pub name: Identifier,
 177     pub values: Vec<String>,
 178 }
 179
 180 impl HasIdentifier for Attribute {
 181     fn identifier(&self) -> &Identifier {
 182         &self.name
 183     }
 184 }
 185
 186 #[derive(Clone, Debug)]
 187 pub struct MultipleResponseSet {
 188     pub name: Identifier,
 189     pub label: String,
 190     pub mr_type: MultipleResponseType,
 191     pub variables: Vec<DictIndex>,
 192 }
 193
 194 impl MultipleResponseSet {
 195     fn with_updated_dict_indexes(
 196         mut self,
 197         f: impl Fn(DictIndex) -> Option<DictIndex>,
 198     ) -> Option<Self> {
 199         update_dict_index_vec(&mut self.variables, f);
 200         (self.variables.len() > 1).then_some(self)
 201     }
 202 }
 203
 204 impl HasIdentifier for MultipleResponseSet {
 205     fn identifier(&self) -> &Identifier {
 206         &self.name
 207     }
 208 }
 209
 210 #[derive(Clone, Debug)]
 211 pub enum MultipleResponseType {
 212     MultipleDichotomy {
 213         value: Value,
 214         labels: CategoryLabels,
 215     },
 216     MultipleCategory,
 217 }
 218
 219 #[derive(Clone, Debug)]
 220 pub struct VariableSet {
 221     pub name: Identifier,
 222     pub variables: Vec<DictIndex>,
 223 }
 224
 225 impl VariableSet {
 226     fn with_updated_dict_indexes(
 227         mut self,
 228         f: impl Fn(DictIndex) -> Option<DictIndex>,
 229     ) -> Option<Self> {
 230         update_dict_index_vec(&mut self.variables, f);
 231         (!self.variables.is_empty()).then_some(self)
 232     }
 233 }
 234
 235 impl HasIdentifier for VariableSet {
 236     fn identifier(&self) -> &Identifier {
 237         &self.name
 238     }
 239 }
 240
 241 #[cfg(test)]
 242 mod test {
 243     use std::collections::HashSet;
 244
 245     use crate::identifier::Identifier;
 246
 247     use super::{ByIdentifier, HasIdentifier};
 248
 249     #[derive(PartialEq, Eq, Debug, Clone)]
 250     struct Variable {
 251         name: Identifier,
 252         value: i32,
 253     }
 254
 255     impl HasIdentifier for Variable {
 256         fn identifier(&self) -> &Identifier {
 257             &self.name
 258         }
 259     }
 260
 261     #[test]
 262     fn test() {
 263         // Variables should not be the same if their values differ.
 264         let abcd = Identifier::new_utf8("abcd").unwrap();
 265         let abcd1 = Variable {
 266             name: abcd.clone(),
 267             value: 1,
 268         };
 269         let abcd2 = Variable {
 270             name: abcd,
 271             value: 2,
 272         };
 273         assert_ne!(abcd1, abcd2);
 274
 275         // But `ByName` should treat them the same.
 276         let abcd1_by_name = ByIdentifier::new(abcd1);
 277         let abcd2_by_name = ByIdentifier::new(abcd2);
 278         assert_eq!(abcd1_by_name, abcd2_by_name);
 279
 280         // And a `HashSet` of `ByName` should also treat them the same.
 281         let mut vars: HashSet<ByIdentifier<Variable>> = HashSet::new();
 282         assert!(vars.insert(ByIdentifier::new(abcd1_by_name.0.clone())));
 283         assert!(!vars.insert(ByIdentifier::new(abcd2_by_name.0.clone())));
 284         assert_eq!(
 285             vars.get(&Identifier::new_utf8("abcd").unwrap())
 286                 .unwrap()
 287                 .0
 288                 .value,
 289             1
 290         );
 291     }
 292 }