1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use crate::{
    context::Context,
    number::{Dimension, NumberParts},
    reply::SearchReply,
};
use std::collections::BinaryHeap;
use std::{
    cmp::{Ord, Ordering, PartialOrd},
    collections::BTreeMap,
};
use strsim::jaro_winkler;

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct SearchResult<'a> {
    score: i32,
    value: &'a str,
}

impl<'a> PartialOrd for SearchResult<'a> {
    fn partial_cmp(&self, other: &SearchResult<'a>) -> Option<Ordering> {
        self.score.partial_cmp(&other.score).map(|x| x.reverse())
    }
}

impl<'a> Ord for SearchResult<'a> {
    fn cmp(&self, other: &SearchResult<'a>) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

pub fn search<'a>(ctx: &'a Context, query: &str, num_results: usize) -> Vec<&'a str> {
    let mut results = BinaryHeap::new();
    let query = query.to_lowercase();
    {
        let mut scan = |x: &'a str| {
            let borrow = x;
            let x = x.to_lowercase();
            let modifier = if x == query {
                4_000
            } else if x.starts_with(&query) {
                3_000
            } else if x.ends_with(&query) {
                2_000
            } else if x.contains(&query) {
                1_000
            } else {
                0_000
            };
            let score = jaro_winkler(&*x, &*query);
            results.push(SearchResult {
                score: (score * 1000.0) as i32 + modifier,
                value: borrow,
            });
            while results.len() > num_results {
                results.pop();
            }
        };

        for k in &ctx.dimensions {
            scan(&**k.id);
        }
        for k in ctx.units.keys() {
            scan(&**k);
        }
        for k in ctx.quantities.values() {
            scan(&**k);
        }
        for k in ctx.substances.keys() {
            scan(&**k);
        }
    }
    results
        .into_sorted_vec()
        .into_iter()
        .filter_map(|x| if x.score > 800 { Some(x.value) } else { None })
        .collect()
}

pub fn query(ctx: &Context, query: &str, num_results: usize) -> SearchReply {
    SearchReply {
        results: search(ctx, query, num_results)
            .into_iter()
            .map(|name| {
                let parts = ctx
                    .lookup(name)
                    .map(|x| x.to_parts(ctx))
                    .or_else(|| {
                        if ctx.substances.get(name).is_some() {
                            Some(NumberParts {
                                quantity: Some("substance".to_owned()),
                                ..Default::default()
                            })
                        } else {
                            None
                        }
                    })
                    .expect("Search returned non-existent result");
                let mut raw = BTreeMap::new();
                raw.insert(Dimension::new(name), 1);
                NumberParts {
                    unit: Some(name.to_owned()),
                    raw_unit: Some(raw),
                    quantity: parts.quantity,
                    ..Default::default()
                }
            })
            .collect(),
    }
}