1use compact_str::CompactString;
4use oxc_allocator::Allocator;
5use oxc_ast::ast::{Comment, Expression, FormalParameters, FunctionBody};
6use oxc_diagnostics::Diagnostics;
7use oxc_span::{SourceType, Span};
8
9use crate::{ParseOptions, Parser, ParserReturn};
10
11mod borrowed;
12mod projection;
13mod shapes;
14pub use borrowed::AdmittedBorrowedExpression;
15pub use projection::{
16 AdmittedExpression, AdmittedHandlerBody, AdmittedParameters, ExpressionObservation,
17 HandlerBodyObservation, ParametersObservation,
18};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum EmbeddingGoal {
22 Expr,
23 HandlerBody,
24 Parameters,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum EmbeddingInputError {
29 SourceTooLarge,
30 UnsupportedProfile,
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum EmbeddingHole {
35 Syntax,
36 UnsupportedFlow,
37 InvalidExpressionShape,
38 InvalidWrappedShape,
39 InvalidModuleContext,
40 InvalidParameterContext,
41}
42
43pub struct EmbeddingInput<'a> {
53 allocator: &'a Allocator,
54 content: &'a str,
55 source_type: SourceType,
56 goal: EmbeddingGoal,
57 wrapper: CompactString,
58 content_span: Span,
59}
60
61impl<'a> EmbeddingInput<'a> {
62 pub fn prepare_in(
65 allocator: &'a Allocator,
66 content: &'a str,
67 source_type: SourceType,
68 goal: EmbeddingGoal,
69 ) -> Result<Self, EmbeddingInputError> {
70 if !source_type.is_module() || source_type.is_typescript_definition() {
71 return Err(EmbeddingInputError::UnsupportedProfile);
72 }
73 let (prefix, suffix) = shapes::wrapper(goal);
74 let length = content
75 .len()
76 .checked_add(prefix.len() + suffix.len())
77 .filter(|length| *length <= crate::MAX_LEN)
78 .ok_or(EmbeddingInputError::SourceTooLarge)?;
79 let content_span = Span::new(prefix.len() as u32, (prefix.len() + content.len()) as u32);
80 let mut wrapper = CompactString::with_capacity(length);
81 wrapper.push_str(prefix);
82 wrapper.push_str(content);
83 wrapper.push_str(suffix);
84 Ok(Self { allocator, content, source_type, goal, wrapper, content_span })
85 }
86
87 #[must_use]
89 pub fn parser_source(&self) -> &str {
90 self.wrapper.as_str()
91 }
92
93 #[must_use]
94 pub const fn parser_content_span(&self) -> Span {
95 self.content_span
96 }
97
98 #[must_use]
109 pub fn observe(self) -> EmbeddingObservation<'a> {
110 let source = self.allocator.alloc_str(self.wrapper.as_str());
111 let options = ParseOptions::default();
112 let parsed =
113 Parser::new(self.allocator, source, self.source_type).with_options(options).parse();
114 let hole = if parsed.is_flow_language {
115 Some(EmbeddingHole::UnsupportedFlow)
116 } else if parsed.panicked || parsed.diagnostics.has_errors() {
117 Some(EmbeddingHole::Syntax)
118 } else {
119 shapes::shape_hole(&parsed.program, self.goal, self.content_span).or_else(|| {
120 (self.goal != EmbeddingGoal::Expr)
121 .then(|| shapes::context_hole(&parsed.program))
122 .flatten()
123 })
124 };
125 EmbeddingObservation {
126 allocator: self.allocator,
127 content: self.content,
128 source_type: self.source_type,
129 goal: self.goal,
130 content_span: self.content_span,
131 parsed,
132 options,
133 hole,
134 }
135 }
136}
137
138pub struct EmbeddingObservation<'a> {
141 allocator: &'a Allocator,
142 content: &'a str,
143 source_type: SourceType,
144 goal: EmbeddingGoal,
145 content_span: Span,
146 parsed: ParserReturn<'a>,
147 options: ParseOptions,
148 hole: Option<EmbeddingHole>,
149}
150
151impl core::fmt::Debug for EmbeddingObservation<'_> {
152 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
153 f.debug_struct("EmbeddingObservation")
154 .field("goal", &self.goal)
155 .field("source_type", &self.source_type)
156 .field("hole", &self.hole)
157 .finish_non_exhaustive()
158 }
159}
160
161impl<'a> EmbeddingObservation<'a> {
162 #[must_use]
163 pub const fn content(&self) -> &'a str {
164 self.content
165 }
166 #[must_use]
167 pub const fn source_type(&self) -> SourceType {
168 self.source_type
169 }
170 #[must_use]
171 pub const fn options(&self) -> ParseOptions {
172 self.options
173 }
174 #[must_use]
175 pub const fn hole(&self) -> Option<EmbeddingHole> {
176 self.hole
177 }
178 #[must_use]
179 pub const fn parser_content_span(&self) -> Span {
180 self.content_span
181 }
182 #[must_use]
183 pub fn comments(&self) -> &[Comment] {
184 self.parsed.program.comments.as_slice()
185 }
186 #[must_use]
187 pub fn diagnostics(&self) -> &Diagnostics {
188 &self.parsed.diagnostics
189 }
190 #[must_use]
191 pub fn panicked(&self) -> bool {
192 self.parsed.panicked
193 }
194 #[must_use]
195 pub fn is_flow_language(&self) -> bool {
196 self.parsed.is_flow_language
197 }
198
199 #[must_use]
200 pub fn expression(&self) -> Option<&Expression<'a>> {
201 (self.hole.is_none() && self.goal == EmbeddingGoal::Expr)
202 .then(|| shapes::expression(&self.parsed.program))
203 .flatten()
204 }
205 #[must_use]
206 pub fn handler_body(&self) -> Option<&FunctionBody<'a>> {
207 (self.hole.is_none() && self.goal == EmbeddingGoal::HandlerBody)
208 .then(|| shapes::arrow(&self.parsed.program).map(|arrow| &*arrow.body))
209 .flatten()
210 }
211 #[must_use]
212 pub fn parameters(&self) -> Option<&FormalParameters<'a>> {
213 (self.hole.is_none() && self.goal == EmbeddingGoal::Parameters)
214 .then(|| shapes::arrow(&self.parsed.program).map(|arrow| &*arrow.params))
215 .flatten()
216 }
217}