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

   1 use std::{
   2     cmp::Ordering,
   3     collections::{HashMap, HashSet},
   4     fmt::Debug,
   5     ops::{Bound, RangeBounds},
   6 };
   7
   8 use encoding_rs::Encoding;
   9 use indexmap::IndexSet;
  10 use num::integer::div_ceil;
  11 use ordered_float::OrderedFloat;
  12
  13 use crate::{
  14     format::Spec,
  15     identifier::{ByIdentifier, HasIdentifier, Identifier},
  16     raw::{self, Alignment, CategoryLabels, Decoder, Measure, MissingValues, RawStr, VarType},
  17 };
  18
  19 pub type DictIndex = usize;
  20
  21 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
  22 pub enum VarWidth {
  23     Numeric,
  24     String(u16),
  25 }
  26
  27 impl PartialOrd for VarWidth {
  28     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
  29         match (self, other) {
  30             (VarWidth::Numeric, VarWidth::Numeric) => Some(Ordering::Equal),
  31             (VarWidth::String(a), VarWidth::String(b)) => Some(a.cmp(b)),
  32             _ => None,
  33         }
  34     }
  35 }
  36
  37 impl VarWidth {
  38     pub const MAX_STRING: u16 = 32767;
  39
  40     pub fn n_dict_indexes(self) -> usize {
  41         match self {
  42             VarWidth::Numeric => 1,
  43             VarWidth::String(w) => div_ceil(w as usize, 8),
  44         }
  45     }
  46
  47     fn width_predicate(
  48         a: Option<VarWidth>,
  49         b: Option<VarWidth>,
  50         f: impl Fn(u16, u16) -> u16,
  51     ) -> Option<VarWidth> {
  52         match (a, b) {
  53             (Some(VarWidth::Numeric), Some(VarWidth::Numeric)) => Some(VarWidth::Numeric),
  54             (Some(VarWidth::String(a)), Some(VarWidth::String(b))) => {
  55                 Some(VarWidth::String(f(a, b)))
  56             }
  57             _ => None,
  58         }
  59     }
  60
  61     /// Returns the wider of `self` and `other`:
  62     /// - Numerical variable widths are equally wide.
  63     /// - Longer strings are wider than shorter strings.
  64     /// - Numerical and string types are incomparable, so result in `None`.
  65     /// - Any `None` in the input yields `None` in the output.
  66     pub fn wider(a: Option<VarWidth>, b: Option<VarWidth>) -> Option<VarWidth> {
  67         Self::width_predicate(a, b, |a, b| a.max(b))
  68     }
  69
  70     /// Returns the narrower of `self` and `other` (see [`Self::wider`]).
  71     pub fn narrower(a: Option<VarWidth>, b: Option<VarWidth>) -> Option<VarWidth> {
  72         Self::width_predicate(a, b, |a, b| a.min(b))
  73     }
  74
  75     pub fn default_display_width(&self) -> u32 {
  76         match self {
  77             VarWidth::Numeric => 8,
  78             VarWidth::String(width) => *width.min(&32) as u32,
  79         }
  80     }
  81 }
  82
  83 impl From<VarWidth> for VarType {
  84     fn from(source: VarWidth) -> Self {
  85         match source {
  86             VarWidth::Numeric => VarType::Numeric,
  87             VarWidth::String(_) => VarType::String,
  88         }
  89     }
  90 }
  91
  92 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
  93 pub enum Value {
  94     Number(Option<OrderedFloat<f64>>),
  95     String(String),
  96 }
  97
  98 impl Value {
  99     pub fn decode(raw: &raw::Value<RawStr<8>>, decoder: &Decoder) -> Self {
 100         match raw {
 101             raw::Value::Number(x) => Value::Number(x.map(|x| x.into())),
 102             raw::Value::String(s) => Value::String(decoder.decode_exact_length(&s.0).into()),
 103         }
 104     }
 105 }
 106
 107 #[derive(Clone, Debug)]
 108 pub struct Dictionary {
 109     pub variables: IndexSet<ByIdentifier<Variable>>,
 110     pub split_file: Vec<DictIndex>,
 111     pub weight: Option<DictIndex>,
 112     pub filter: Option<DictIndex>,
 113     pub case_limit: Option<u64>,
 114     pub file_label: Option<String>,
 115     pub documents: Vec<String>,
 116     pub vectors: HashSet<ByIdentifier<Vector>>,
 117     pub attributes: HashMap<Identifier, Vec<String>>,
 118     pub mrsets: HashSet<ByIdentifier<MultipleResponseSet>>,
 119     pub variable_sets: HashSet<ByIdentifier<VariableSet>>,
 120     pub encoding: &'static Encoding,
 121 }
 122
 123 pub struct DuplicateVariableName;
 124
 125 impl Dictionary {
 126     pub fn new(encoding: &'static Encoding) -> Self {
 127         Self {
 128             variables: IndexSet::new(),
 129             split_file: Vec::new(),
 130             weight: None,
 131             filter: None,
 132             case_limit: None,
 133             file_label: None,
 134             documents: Vec::new(),
 135             vectors: HashSet::new(),
 136             attributes: HashMap::new(),
 137             mrsets: HashSet::new(),
 138             variable_sets: HashSet::new(),
 139             encoding,
 140         }
 141     }
 142
 143     pub fn add_var(&mut self, variable: Variable) -> Result<(), DuplicateVariableName> {
 144         if self.variables.insert(ByIdentifier::new(variable)) {
 145             Ok(())
 146         } else {
 147             Err(DuplicateVariableName)
 148         }
 149     }
 150
 151     pub fn reorder_var(&mut self, from_index: DictIndex, to_index: DictIndex) {
 152         if from_index != to_index {
 153             self.variables.move_index(from_index, to_index);
 154             self.update_dict_indexes(&|index| {
 155                 #[allow(clippy::collapsible_else_if)]
 156                 if index == from_index {
 157                     Some(to_index)
 158                 } else if from_index < to_index {
 159                     if index > from_index && index <= to_index {
 160                         Some(index - 1)
 161                     } else {
 162                         Some(index)
 163                     }
 164                 } else {
 165                     if index >= to_index && index < from_index {
 166                         Some(index + 1)
 167                     } else {
 168                         Some(index)
 169                     }
 170                 }
 171             })
 172         }
 173     }
 174
 175     pub fn retain_vars<F>(&mut self, keep: F)
 176     where
 177         F: Fn(&Variable) -> bool,
 178     {
 179         let mut deleted = Vec::new();
 180         let mut index = 0;
 181         self.variables.retain(|var_by_id| {
 182             let keep = keep(&var_by_id.0);
 183             if !keep {
 184                 deleted.push(index);
 185             }
 186             index += 1;
 187             keep
 188         });
 189         if !deleted.is_empty() {
 190             self.update_dict_indexes(&|index| match deleted.binary_search(&index) {
 191                 Ok(_) => None,
 192                 Err(position) => Some(position),
 193             })
 194         }
 195     }
 196
 197     pub fn delete_vars<R>(&mut self, range: R)
 198     where
 199         R: RangeBounds<DictIndex>,
 200     {
 201         let start = match range.start_bound() {
 202             Bound::Included(&start) => start,
 203             Bound::Excluded(&start) => start + 1,
 204             Bound::Unbounded => 0,
 205         };
 206         let end = match range.end_bound() {
 207             Bound::Included(&end) => end + 1,
 208             Bound::Excluded(&end) => end,
 209             Bound::Unbounded => self.variables.len(),
 210         };
 211         if end > start {
 212             self.variables.drain(start..end);
 213             self.update_dict_indexes(&|index| {
 214                 if index < start {
 215                     Some(index)
 216                 } else if index < end {
 217                     None
 218                 } else {
 219                     Some(index - end - start)
 220                 }
 221             })
 222         }
 223     }
 224
 225     fn update_dict_indexes<F>(&mut self, f: &F)
 226     where
 227         F: Fn(DictIndex) -> Option<DictIndex>,
 228     {
 229         update_dict_index_vec(&mut self.split_file, f);
 230         self.weight = self.weight.and_then(f);
 231         self.filter = self.filter.and_then(f);
 232         self.vectors = self
 233             .vectors
 234             .drain()
 235             .filter_map(|vector_by_id| {
 236                 vector_by_id
 237                     .0
 238                     .with_updated_dict_indexes(f)
 239                     .map(ByIdentifier::new)
 240             })
 241             .collect();
 242         self.mrsets = self
 243             .mrsets
 244             .drain()
 245             .filter_map(|mrset_by_id| {
 246                 mrset_by_id
 247                     .0
 248                     .with_updated_dict_indexes(f)
 249                     .map(ByIdentifier::new)
 250             })
 251             .collect();
 252         self.variable_sets = self
 253             .variable_sets
 254             .drain()
 255             .filter_map(|var_set_by_id| {
 256                 var_set_by_id
 257                     .0
 258                     .with_updated_dict_indexes(f)
 259                     .map(ByIdentifier::new)
 260             })
 261             .collect();
 262     }
 263 }
 264
 265 fn update_dict_index_vec<F>(dict_indexes: &mut Vec<DictIndex>, f: F)
 266 where
 267     F: Fn(DictIndex) -> Option<DictIndex>,
 268 {
 269     dict_indexes.retain_mut(|index| {
 270         if let Some(new) = f(*index) {
 271             *index = new;
 272             true
 273         } else {
 274             false
 275         }
 276     });
 277 }
 278
 279 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 280 pub enum Role {
 281     Input,
 282     Target,
 283     Both,
 284     None,
 285     Partition,
 286     Split,
 287 }
 288
 289 impl Default for Role {
 290     fn default() -> Self {
 291         Self::Input
 292     }
 293 }
 294
 295 pub enum DictClass {
 296     Ordinary,
 297     System,
 298     Scratch,
 299 }
 300
 301 impl DictClass {
 302     pub fn from_identifier(id: &Identifier) -> Self {
 303         if id.0.starts_with('$') {
 304             Self::System
 305         } else if id.0.starts_with('#') {
 306             Self::Scratch
 307         } else {
 308             Self::Ordinary
 309         }
 310     }
 311
 312     pub fn must_leave(self) -> bool {
 313         match self {
 314             DictClass::Ordinary => false,
 315             DictClass::System => false,
 316             DictClass::Scratch => true,
 317         }
 318     }
 319 }
 320
 321 #[derive(Clone, Debug)]
 322 pub struct Variable {
 323     pub name: Identifier,
 324     pub width: VarWidth,
 325     pub missing_values: MissingValues,
 326     pub print_format: Spec,
 327     pub write_format: Spec,
 328     pub value_labels: HashMap<Value, String>,
 329     pub label: Option<String>,
 330     pub measure: Option<Measure>,
 331     pub role: Role,
 332     pub display_width: u32,
 333     pub alignment: Alignment,
 334     pub leave: bool,
 335     pub short_names: Vec<Identifier>,
 336     pub attributes: HashSet<ByIdentifier<Attribute>>,
 337 }
 338
 339 impl Variable {
 340     pub fn new(name: Identifier, width: VarWidth) -> Self {
 341         let var_type = VarType::from_width(width);
 342         let leave = DictClass::from_identifier(&name).must_leave();
 343         Self {
 344             name,
 345             width,
 346             missing_values: MissingValues::default(),
 347             print_format: Spec::default_for_width(width),
 348             write_format: Spec::default_for_width(width),
 349             value_labels: HashMap::new(),
 350             label: None,
 351             measure: Measure::default_for_type(var_type),
 352             role: Role::default(),
 353             display_width: width.default_display_width(),
 354             alignment: Alignment::default_for_type(var_type),
 355             leave,
 356             short_names: Vec::new(),
 357             attributes: HashSet::new(),
 358         }
 359     }
 360 }
 361
 362 impl HasIdentifier for Variable {
 363     fn identifier(&self) -> &Identifier {
 364         &self.name
 365     }
 366 }
 367
 368 #[derive(Clone, Debug)]
 369 pub struct Vector {
 370     pub name: Identifier,
 371     pub variables: Vec<DictIndex>,
 372 }
 373
 374 impl Vector {
 375     fn with_updated_dict_indexes(
 376         mut self,
 377         f: impl Fn(DictIndex) -> Option<DictIndex>,
 378     ) -> Option<Self> {
 379         update_dict_index_vec(&mut self.variables, f);
 380         (!self.variables.is_empty()).then_some(self)
 381     }
 382 }
 383
 384 impl HasIdentifier for Vector {
 385     fn identifier(&self) -> &Identifier {
 386         &self.name
 387     }
 388 }
 389
 390 #[derive(Clone, Debug)]
 391 pub struct Attribute {
 392     pub name: Identifier,
 393     pub values: Vec<String>,
 394 }
 395
 396 impl HasIdentifier for Attribute {
 397     fn identifier(&self) -> &Identifier {
 398         &self.name
 399     }
 400 }
 401
 402 #[derive(Clone, Debug)]
 403 pub struct MultipleResponseSet {
 404     pub name: Identifier,
 405     pub label: String,
 406     pub mr_type: MultipleResponseType,
 407     pub variables: Vec<DictIndex>,
 408 }
 409
 410 impl MultipleResponseSet {
 411     fn with_updated_dict_indexes(
 412         mut self,
 413         f: impl Fn(DictIndex) -> Option<DictIndex>,
 414     ) -> Option<Self> {
 415         update_dict_index_vec(&mut self.variables, f);
 416         (self.variables.len() > 1).then_some(self)
 417     }
 418 }
 419
 420 impl HasIdentifier for MultipleResponseSet {
 421     fn identifier(&self) -> &Identifier {
 422         &self.name
 423     }
 424 }
 425
 426 #[derive(Clone, Debug)]
 427 pub enum MultipleResponseType {
 428     MultipleDichotomy {
 429         value: Value,
 430         labels: CategoryLabels,
 431     },
 432     MultipleCategory,
 433 }
 434
 435 #[derive(Clone, Debug)]
 436 pub struct VariableSet {
 437     pub name: Identifier,
 438     pub variables: Vec<DictIndex>,
 439 }
 440
 441 impl VariableSet {
 442     fn with_updated_dict_indexes(
 443         mut self,
 444         f: impl Fn(DictIndex) -> Option<DictIndex>,
 445     ) -> Option<Self> {
 446         update_dict_index_vec(&mut self.variables, f);
 447         (!self.variables.is_empty()).then_some(self)
 448     }
 449 }
 450
 451 impl HasIdentifier for VariableSet {
 452     fn identifier(&self) -> &Identifier {
 453         &self.name
 454     }
 455 }
 456
 457 #[cfg(test)]
 458 mod test {
 459     use std::collections::HashSet;
 460
 461     use crate::identifier::Identifier;
 462
 463     use super::{ByIdentifier, HasIdentifier};
 464
 465     #[derive(PartialEq, Eq, Debug, Clone)]
 466     struct Variable {
 467         name: Identifier,
 468         value: i32,
 469     }
 470
 471     impl HasIdentifier for Variable {
 472         fn identifier(&self) -> &Identifier {
 473             &self.name
 474         }
 475     }
 476
 477     #[test]
 478     fn test() {
 479         // Variables should not be the same if their values differ.
 480         let abcd = Identifier::new_utf8("abcd").unwrap();
 481         let abcd1 = Variable {
 482             name: abcd.clone(),
 483             value: 1,
 484         };
 485         let abcd2 = Variable {
 486             name: abcd,
 487             value: 2,
 488         };
 489         assert_ne!(abcd1, abcd2);
 490
 491         // But `ByName` should treat them the same.
 492         let abcd1_by_name = ByIdentifier::new(abcd1);
 493         let abcd2_by_name = ByIdentifier::new(abcd2);
 494         assert_eq!(abcd1_by_name, abcd2_by_name);
 495
 496         // And a `HashSet` of `ByName` should also treat them the same.
 497         let mut vars: HashSet<ByIdentifier<Variable>> = HashSet::new();
 498         assert!(vars.insert(ByIdentifier::new(abcd1_by_name.0.clone())));
 499         assert!(!vars.insert(ByIdentifier::new(abcd2_by_name.0.clone())));
 500         assert_eq!(
 501             vars.get(&Identifier::new_utf8("abcd").unwrap())
 502                 .unwrap()
 503                 .0
 504                 .value,
 505             1
 506         );
 507     }
 508 }