1use crate::{
7 dispatcher::PROMPT_ARROW,
8 prelude::{COMMAND_LEADER, ChiselCommand, PROMPT_ARROW_STR},
9};
10use rustyline::{
11 Helper,
12 completion::Completer,
13 highlight::{CmdKind, Highlighter},
14 hint::Hinter,
15 validate::{ValidationContext, ValidationResult, Validator},
16};
17use solar::parse::{
18 Lexer,
19 interface::Session,
20 token::{Token, TokenKind},
21};
22use std::{borrow::Cow, ops::Range, str::FromStr};
23use yansi::{Color, Style};
24
25const MAX_ANSI_LEN: usize = 9;
33
34pub struct SolidityHelper {
36 errored: bool,
37
38 do_paint: bool,
39 sess: Session,
40}
41
42impl Default for SolidityHelper {
43 fn default() -> Self {
44 Self::new()
45 }
46}
47
48impl SolidityHelper {
49 pub fn new() -> Self {
51 Self {
52 errored: false,
53 do_paint: yansi::is_enabled(),
54 sess: Session::builder().with_silent_emitter(None).build(),
55 }
56 }
57
58 pub fn errored(&self) -> bool {
60 self.errored
61 }
62
63 pub fn set_errored(&mut self, errored: bool) -> &mut Self {
65 self.errored = errored;
66 self
67 }
68
69 pub fn highlight<'a>(&self, input: &'a str) -> Cow<'a, str> {
71 if !self.do_paint() {
72 return Cow::Borrowed(input);
73 }
74
75 if input.starts_with(COMMAND_LEADER) {
77 let (cmd, rest) = match input.split_once(' ') {
78 Some((cmd, rest)) => (cmd, Some(rest)),
79 None => (input, None),
80 };
81 let cmd = cmd.strip_prefix(COMMAND_LEADER).unwrap_or(cmd);
82
83 let mut out = String::with_capacity(input.len() + MAX_ANSI_LEN);
84
85 out.push(COMMAND_LEADER);
87 let cmd_res = ChiselCommand::from_str(cmd);
88 let style = (if cmd_res.is_ok() { Color::Green } else { Color::Red }).foreground();
89 Self::paint_unchecked(cmd, style, &mut out);
90
91 match rest {
93 Some(rest) if !rest.is_empty() => {
94 out.push(' ');
95 out.push_str(rest);
96 }
97 _ => {}
98 }
99
100 Cow::Owned(out)
101 } else {
102 let mut out = String::with_capacity(input.len() * 2);
103 self.with_contiguous_styles(input, |style, range| {
104 Self::paint_unchecked(&input[range], style, &mut out);
105 });
106 Cow::Owned(out)
107 }
108 }
109
110 fn with_contiguous_styles(&self, input: &str, mut f: impl FnMut(Style, Range<usize>)) {
114 self.enter(|sess| {
115 let len = input.len();
116 let mut index = 0;
117 for token in Lexer::new(sess, input) {
118 let range = token.span.lo().to_usize()..token.span.hi().to_usize();
119 let style = token_style(&token);
120 if index < range.start {
121 f(Style::default(), index..range.start);
122 }
123 index = range.end;
124 f(style, range);
125 }
126 if index < len {
127 f(Style::default(), index..len);
128 }
129 });
130 }
131
132 fn validate_closed(&self, input: &str) -> ValidationResult {
134 let mut depth = [0usize; 3];
135 self.enter(|sess| {
136 for token in Lexer::new(sess, input) {
137 match token.kind {
138 TokenKind::OpenDelim(delim) => {
139 depth[delim as usize] += 1;
140 }
141 TokenKind::CloseDelim(delim) => {
142 depth[delim as usize] = depth[delim as usize].saturating_sub(1);
143 }
144 _ => {}
145 }
146 }
147 });
148 if depth == [0; 3] { ValidationResult::Valid(None) } else { ValidationResult::Incomplete }
149 }
150
151 fn paint_unchecked(string: &str, style: Style, out: &mut String) {
154 if style == Style::default() {
155 out.push_str(string);
156 } else {
157 let _ = style.fmt_prefix(out);
158 out.push_str(string);
159 let _ = style.fmt_suffix(out);
160 }
161 }
162
163 fn paint_unchecked_owned(string: &str, style: Style) -> String {
164 let mut out = String::with_capacity(MAX_ANSI_LEN + string.len());
165 Self::paint_unchecked(string, style, &mut out);
166 out
167 }
168
169 fn do_paint(&self) -> bool {
171 self.do_paint
172 }
173
174 fn enter(&self, f: impl FnOnce(&Session)) {
176 self.sess.enter_sequential(|| f(&self.sess));
177 }
178}
179
180impl Highlighter for SolidityHelper {
181 fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
182 self.highlight(line)
183 }
184
185 fn highlight_char(&self, line: &str, pos: usize, _kind: CmdKind) -> bool {
186 pos == line.len()
187 }
188
189 fn highlight_prompt<'b, 's: 'b, 'p: 'b>(
190 &'s self,
191 prompt: &'p str,
192 _default: bool,
193 ) -> Cow<'b, str> {
194 if !self.do_paint() {
195 return Cow::Borrowed(prompt);
196 }
197
198 let mut out = prompt.to_string();
199
200 if prompt.starts_with("(ID: ") {
202 let id_end = prompt.find(')').unwrap();
203 let id_span = 5..id_end;
204 let id = &prompt[id_span.clone()];
205 out.replace_range(
206 id_span,
207 &Self::paint_unchecked_owned(id, Color::Yellow.foreground()),
208 );
209 out.replace_range(1..=2, &Self::paint_unchecked_owned("ID", Color::Cyan.foreground()));
210 }
211
212 if let Some(i) = out.find(PROMPT_ARROW) {
213 let style =
214 if self.errored { Color::Red.foreground() } else { Color::Green.foreground() };
215 out.replace_range(i..=i + 2, &Self::paint_unchecked_owned(PROMPT_ARROW_STR, style));
216 }
217
218 Cow::Owned(out)
219 }
220}
221
222impl Validator for SolidityHelper {
223 fn validate(&self, ctx: &mut ValidationContext<'_>) -> rustyline::Result<ValidationResult> {
224 Ok(self.validate_closed(ctx.input()))
225 }
226}
227
228impl Completer for SolidityHelper {
229 type Candidate = String;
230}
231
232impl Hinter for SolidityHelper {
233 type Hint = String;
234}
235
236impl Helper for SolidityHelper {}
237
238#[expect(non_upper_case_globals)]
239#[deny(unreachable_patterns)]
240fn token_style(token: &Token) -> Style {
241 use solar::parse::{
242 interface::kw::*,
243 token::{TokenKind::*, TokenLitKind::*},
244 };
245
246 match token.kind {
247 Literal(Str | HexStr | UnicodeStr, _) => Color::Green.foreground(),
248 Literal(..) => Color::Yellow.foreground(),
249
250 Ident(
251 Memory | Storage | Calldata | Public | Private | Internal | External | Constant | Pure
252 | View | Payable | Anonymous | Indexed | Abstract | Virtual | Override | Modifier
253 | Immutable | Unchecked,
254 ) => Color::Cyan.foreground(),
255
256 Ident(s) if s.is_elementary_type() => Color::Blue.foreground(),
257 Ident(Mapping) => Color::Blue.foreground(),
258
259 Ident(s) if s.is_used_keyword() || s.is_yul_keyword() => Color::Magenta.foreground(),
260 Arrow | FatArrow => Color::Magenta.foreground(),
261
262 Comment(..) => Color::Primary.dim(),
263
264 _ => Color::Primary.foreground(),
265 }
266}