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