gravel_core/
engine.rs

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
use crate::{performance::Stopwatch, scoring, timed, CoreMessage};
use abi_stable::std_types::RString;
use abi_stable::{external_types::crossbeam_channel::RSender, sabi_trait, std_types::RStr, traits::IntoReprRust};
use gravel_ffi::{ActionKind, ArcDynHit, FrontendMessage, HitActionContext, ProviderResult, RefDynHitActionContext};
use gravel_ffi::{BoxDynProvider, QueryResult};
use itertools::Itertools;
use std::iter::once;

/// Holds a [`BoxDynProvider`] and some additional metadata.
struct ProviderInfo {
	pub name: String,
	pub provider: BoxDynProvider,
	pub keyword: Option<String>,
}

/// Aggregates and scores hits from the given [`BoxDynProvider`]s.
pub struct QueryEngine {
	providers: Vec<ProviderInfo>,
	action_context: ActionContext,
}

impl QueryEngine {
	pub fn new(sender: RSender<CoreMessage>) -> Self {
		Self {
			providers: vec![],
			action_context: ActionContext::new(sender),
		}
	}

	/// Adds the provider to the engine's collection.
	pub fn register(&mut self, name: String, provider: BoxDynProvider, keyword: Option<String>) -> &mut Self {
		let info = ProviderInfo {
			name,
			provider,
			keyword,
		};

		self.providers.push(info);
		self
	}

	pub fn query(&self, query: RStr<'_>) -> QueryResult {
		fn inner(engine: &QueryEngine, query: &str) -> QueryResult {
			if let Some(result) = engine.try_keyword_query(query) {
				return result;
			}

			engine.full_query(query)
		}

		let query = query.into_rust();

		if query.trim().is_empty() {
			return QueryResult::default();
		}

		inner(self, query)
	}

	pub fn run_hit_action(&self, hit: &ArcDynHit, kind: ActionKind) {
		let context = (&self.action_context).into();

		match kind {
			ActionKind::Primary => hit.action(context),
			ActionKind::Secondary => hit.secondary_action(context),
		};
	}

	pub fn clear_caches(&self) {
		for provider in &self.providers {
			provider.provider.clear_caches();
		}
	}

	/// Runs the query against all available providers.
	fn full_query(&self, query: &str) -> QueryResult {
		let providers = self.providers.iter().filter(|provider| provider.keyword.is_none());

		query_all(providers, query)
	}

	/// Tries to find a provider with the a keyword that matches the query's.
	/// If one is found, the keyword is stripped from the query and the
	/// resulting new query is ran against that provider only.
	fn try_keyword_query(&self, query: &str) -> Option<QueryResult> {
		let first_word = query.split(' ').next()?;

		let provider = self.check_keywords(first_word)?;

		// remove the keyword from the query
		let new_query = &query[first_word.len()..query.len()].trim_start();

		Some(query_all(once(provider), new_query))
	}

	/// Tries to find a provider with the a keyword that matches the given string.
	fn check_keywords(&self, first_word: &str) -> Option<&ProviderInfo> {
		self.providers
			.iter()
			.find(|p| matches!(&p.keyword, Some(k) if k == first_word))
	}
}

/// Queries providers; aggregates, scores and orders [`ArcDynHit`]s.
#[allow(single_use_lifetimes)]
fn query_all<'a>(providers: impl Iterator<Item = &'a ProviderInfo>, query: &str) -> QueryResult {
	let hits = providers.flat_map(|p| query_one(p, query).hits).collect_vec();

	timed!("scoring took", {
		let hits = match query.trim() {
			"*" => scoring::to_unscored(hits),
			_ => scoring::to_scored(hits, query),
		};

		QueryResult::new(hits)
	})
}

fn query_one(info: &ProviderInfo, query: &str) -> ProviderResult {
	let stopwatch = Stopwatch::start();

	let result = info.provider.query(query.into());

	log::trace!(
		"query for provider '{}' took {stopwatch} and produced {} hits",
		info.name,
		result.hits.len()
	);

	result
}

struct ActionContext {
	sender: RSender<CoreMessage>,
}

impl ActionContext {
	pub fn new(sender: RSender<CoreMessage>) -> Self {
		Self { sender }
	}

	fn send(&self, message: CoreMessage) {
		self.sender
			.send(message)
			.inspect_err(|e| log::error!("unable to send core message: {e}"))
			.ok();
	}

	fn send_frontend(&self, message: FrontendMessage) {
		self.send(CoreMessage::Frontend(message));
	}
}

impl<'a> From<&'a ActionContext> for RefDynHitActionContext<'a> {
	fn from(value: &'a ActionContext) -> Self {
		Self::from_ptr(value, sabi_trait::TD_Opaque)
	}
}

impl HitActionContext for ActionContext {
	fn hide_frontend(&self) {
		self.send_frontend(FrontendMessage::Hide);
	}

	fn refresh_frontend(&self) {
		self.send_frontend(FrontendMessage::Refresh);
	}

	fn exit(&self) {
		self.send_frontend(FrontendMessage::Exit);
	}

	fn restart(&self) {
		self.send_frontend(FrontendMessage::Restart);
	}

	fn set_query(&self, query: RString) {
		self.send_frontend(FrontendMessage::ShowWithQuery(query));
	}

	fn clear_caches(&self) {
		self.send(CoreMessage::ClearCaches);
	}
}