Skip to main content

mago_analyzer/external/
scan.rs

1use std::sync::Arc;
2use std::time::Instant;
3
4use mago_database::GlobSettings;
5use mago_database::file::File;
6use mago_database::file::FileId;
7use mago_database::file::FileType;
8use mago_database::matcher::ExclusionMatcher;
9use mago_extension::PayloadWriter;
10use mago_extension::source::SourceSnapshot;
11use mago_names::ResolvedNames;
12use mago_syntax::cst::Program;
13
14use crate::external::AnalyzerTransport;
15use crate::external::Backend;
16use crate::external::ExternalAnalyzerError;
17use crate::external::error::protocol;
18use crate::external::protocol;
19
20const BOOTSTRAP_GROUP: u64 = 0x434F_4445_5343_414E;
21const MESSAGE_OVERHEAD: usize = 12 + 1 + 1 + 4 + 4;
22
23#[derive(Debug, Clone)]
24struct HookMatcher {
25    index: u16,
26    paths: ExclusionMatcher<String>,
27}
28
29#[derive(Debug, Clone)]
30struct BackendPlan {
31    backend: u16,
32    hooks: Box<[HookMatcher]>,
33}
34
35/// Compiled source-file selectors advertised by enabled external plugins.
36#[derive(Debug, Clone)]
37pub struct CodebaseScanPlan {
38    backends: Arc<[BackendPlan]>,
39}
40
41impl CodebaseScanPlan {
42    pub(super) fn compile<T>(backends: &[Backend<T>]) -> Result<Option<Self>, ExternalAnalyzerError> {
43        let mut plans = Vec::new();
44        let mut hook_count = 0usize;
45        let mut target_count = 0usize;
46        for (backend, registered) in backends.iter().enumerate() {
47            let mut hooks = Vec::with_capacity(registered.registration.codebase_scan_hooks.len());
48            for hook in &registered.registration.codebase_scan_hooks {
49                hook_count += 1;
50                target_count += hook.targets.len();
51                let paths = ExclusionMatcher::compile(hook.targets.iter().cloned(), GlobSettings::default()).map_err(
52                    |error| {
53                        protocol(format!(
54                            "codebase-scan hook {} contains an invalid source-file target: {error}",
55                            hook.index
56                        ))
57                    },
58                )?;
59                hooks.push(HookMatcher { index: hook.index, paths });
60            }
61
62            if !hooks.is_empty() {
63                let backend = u16::try_from(backend)
64                    .map_err(|_| protocol("more than 65,536 external analyzer backends were configured"))?;
65                plans.push(BackendPlan { backend, hooks: hooks.into_boxed_slice() });
66            }
67        }
68
69        tracing::trace!(
70            backends = plans.len(),
71            hooks = hook_count,
72            targets = target_count,
73            "Compiled external codebase-scan selectors."
74        );
75        Ok((!plans.is_empty()).then(|| Self { backends: plans.into() }))
76    }
77
78    /// Captures one matching host file while its first parsed syntax tree is live.
79    ///
80    /// Returns `None` without walking or encoding the tree when no hook target
81    /// matches the file's logical path.
82    ///
83    /// # Errors
84    ///
85    /// Returns an error when the syntax snapshot exceeds protocol limits.
86    pub fn capture(
87        &self,
88        file: &Arc<File>,
89        program: &Program<'_>,
90        resolved_names: &ResolvedNames<'_>,
91    ) -> Result<Option<CodebaseScanFile>, ExternalAnalyzerError> {
92        if file.file_type != FileType::Host {
93            return Ok(None);
94        }
95
96        let Ok(path) = std::str::from_utf8(&file.name) else {
97            return Ok(None);
98        };
99        let mut routes = Vec::new();
100        for backend in self.backends.iter() {
101            let hooks = backend
102                .hooks
103                .iter()
104                .filter_map(|hook| hook.paths.is_match(path).then_some(hook.index))
105                .collect::<Vec<_>>();
106            if !hooks.is_empty() {
107                routes.push(CodebaseScanRoute { backend: backend.backend, hooks: hooks.into_boxed_slice() });
108            }
109        }
110
111        if routes.is_empty() {
112            return Ok(None);
113        }
114
115        let snapshot = SourceSnapshot::complete_with_literals(program, resolved_names)?;
116        let mut writer = PayloadWriter::with_capacity(snapshot.encoded_len_with_literals());
117        snapshot.write_to_with_literals(&mut writer)?;
118        Ok(Some(CodebaseScanFile {
119            file: Arc::clone(file),
120            snapshot: writer.finish().into(),
121            routes: routes.into_boxed_slice(),
122        }))
123    }
124}
125
126#[derive(Debug, Clone)]
127struct CodebaseScanRoute {
128    backend: u16,
129    hooks: Box<[u16]>,
130}
131
132/// An owned selected-source snapshot that can outlive its parser arena.
133#[derive(Debug, Clone)]
134pub struct CodebaseScanFile {
135    file: Arc<File>,
136    snapshot: Arc<[u8]>,
137    routes: Box<[CodebaseScanRoute]>,
138}
139
140impl CodebaseScanFile {
141    #[inline]
142    #[must_use]
143    pub fn file_id(&self) -> FileId {
144        self.file.id
145    }
146
147    fn route(&self, backend: u16) -> Option<&[u16]> {
148        self.routes.iter().find(|route| route.backend == backend).map(|route| route.hooks.as_ref())
149    }
150
151    fn encoded_len(&self, hooks: &[u16]) -> usize {
152        4usize
153            .saturating_add(hooks.len().saturating_mul(2))
154            .saturating_add(4)
155            .saturating_add(self.file.name.len())
156            .saturating_add(4)
157            .saturating_add(self.file.contents.len())
158            .saturating_add(self.snapshot.len())
159    }
160}
161
162pub(super) fn dispatch<T>(
163    backends: &[Backend<T>],
164    mut files: Vec<CodebaseScanFile>,
165) -> Result<(), ExternalAnalyzerError>
166where
167    T: AnalyzerTransport,
168{
169    files.sort_unstable_by(|left, right| left.file.name.cmp(&right.file.name));
170    for (backend_index, backend) in backends.iter().enumerate() {
171        if backend.registration.codebase_scan_hooks.is_empty() {
172            continue;
173        }
174
175        let backend_index = u16::try_from(backend_index)
176            .map_err(|_| protocol("more than 65,536 external analyzer backends were configured"))?;
177        let selected =
178            files.iter().filter_map(|file| file.route(backend_index).map(|hooks| (file, hooks))).collect::<Vec<_>>();
179        let active_hooks = backend.registration.codebase_scan_hooks.iter().map(|hook| hook.index).collect::<Vec<_>>();
180        let maximum = backend.transport.maximum_payload_size();
181        let batches = encode_batches(&active_hooks, &selected, maximum)?;
182        let request_bytes = batches.iter().map(Vec::len).sum::<usize>();
183        let started_at = tracing::enabled!(tracing::Level::TRACE).then(Instant::now);
184        tracing::trace!(
185            backend = backend_index,
186            files = selected.len(),
187            batches = batches.len(),
188            request_bytes,
189            "Broadcasting filtered codebase-scan snapshots."
190        );
191        let responses = backend.transport.broadcast_sequence(BOOTSTRAP_GROUP, &batches)?;
192        for batch in responses {
193            for response in batch {
194                protocol::decode_codebase_scan_response(&response)?;
195            }
196        }
197        if let Some(started_at) = started_at {
198            tracing::trace!(
199                backend = backend_index,
200                files = selected.len(),
201                batches = batches.len(),
202                request_bytes,
203                elapsed = ?started_at.elapsed(),
204                "Filtered codebase-scan broadcast completed."
205            );
206        }
207    }
208
209    Ok(())
210}
211
212fn encode_batches(
213    active_hooks: &[u16],
214    files: &[(&CodebaseScanFile, &[u16])],
215    maximum: usize,
216) -> Result<Vec<Vec<u8>>, ExternalAnalyzerError> {
217    let message_overhead = MESSAGE_OVERHEAD.saturating_add(active_hooks.len().saturating_mul(2));
218    let mut ranges = Vec::new();
219    let mut start = 0;
220    let mut length = message_overhead;
221    for (index, (file, hooks)) in files.iter().enumerate() {
222        let record = file.encoded_len(hooks);
223        if message_overhead.saturating_add(record) > maximum {
224            return Err(protocol(format!(
225                "codebase-scan snapshot for `{}` requires {} bytes, exceeding the worker payload limit of {maximum}",
226                mago_bytes::BytesDisplay(&file.file.name),
227                message_overhead.saturating_add(record),
228            )));
229        }
230        if index > start && length.saturating_add(record) > maximum {
231            ranges.push(start..index);
232            start = index;
233            length = message_overhead;
234        }
235        length = length.saturating_add(record);
236    }
237    ranges.push(start..files.len());
238
239    let range_count = ranges.len();
240    ranges
241        .into_iter()
242        .enumerate()
243        .map(|(index, range)| {
244            let mut writer = protocol::message_writer_with_capacity(
245                protocol::CODEBASE_SCAN_REQUEST,
246                files[range.clone()]
247                    .iter()
248                    .fold(message_overhead, |size, (file, hooks)| size.saturating_add(file.encoded_len(hooks))),
249            );
250            writer.write_bool(index == 0);
251            writer.write_bool(index + 1 == range_count);
252            writer.write_length(active_hooks.len())?;
253            for hook in active_hooks {
254                writer.write_u16(*hook);
255            }
256            writer.write_length(range.len())?;
257            for (file, hooks) in &files[range] {
258                writer.write_length(hooks.len())?;
259                for hook in *hooks {
260                    writer.write_u16(*hook);
261                }
262                writer.write_bytes(&file.file.name)?;
263                writer.write_bytes(&file.file.contents)?;
264                writer.write_raw(&file.snapshot);
265            }
266            Ok(writer.finish())
267        })
268        .collect()
269}