difftreelog
refactor turn Thunk structure around to reduce allocations
in: master
9 files changed
.gitignorediffbeforeafterboth--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,10 @@
.vscode
.direnv
+# Nix artifacts
+/result
+/result-*
+
cache
jsonnet-cpp
crates/jrsonnet-evaluator/src/arr/spec.rsdiffbeforeafterboth--- a/crates/jrsonnet-evaluator/src/arr/spec.rs
+++ b/crates/jrsonnet-evaluator/src/arr/spec.rs
@@ -7,7 +7,7 @@
use super::ArrValue;
use crate::{
error::ErrorKind::InfiniteRecursionDetected, evaluate, function::FuncVal, typed::Typed,
- Context, Error, ObjValue, Result, Thunk, Val,
+ val::ThunkValue, Context, Error, ObjValue, Result, Thunk, Val,
};
pub trait ArrayLike: Any + Trace + Debug {
@@ -191,9 +191,25 @@
ArrayThunk::Waiting(_) | ArrayThunk::Pending => {}
};
- let arr_thunk = self.clone();
- Some(Thunk!(move || {
- arr_thunk.get(index).transpose().expect("index checked")
+ #[derive(Trace)]
+ struct ExprArrThunk {
+ expr: ExprArray,
+ index: usize,
+ }
+ impl ThunkValue for ExprArrThunk {
+ type Output = Val;
+
+ fn get(&self) -> Result<Self::Output> {
+ self.expr
+ .get(self.index)
+ .transpose()
+ .expect("index checked")
+ }
+ }
+
+ Some(Thunk::new(ExprArrThunk {
+ expr: self.clone(),
+ index,
}))
}
fn get_cheap(&self, _index: usize) -> Option<Val> {
@@ -484,9 +500,22 @@
ArrayThunk::Waiting(()) | ArrayThunk::Pending => {}
};
- let arr_thunk = self.clone();
- Some(Thunk!(move || {
- arr_thunk.get(index).transpose().expect("index checked")
+ #[derive(Trace)]
+ struct MappedArrayThunk<const WITH_INDEX: bool> {
+ arr: MappedArray<WITH_INDEX>,
+ index: usize,
+ }
+ impl<const WITH_INDEX: bool> ThunkValue for MappedArrayThunk<WITH_INDEX> {
+ type Output = Val;
+
+ fn get(&self) -> Result<Self::Output> {
+ self.arr.get(self.index).transpose().expect("index checked")
+ }
+ }
+
+ Some(Thunk::new(MappedArrayThunk {
+ arr: self.clone(),
+ index,
}))
}
crates/jrsonnet-evaluator/src/dynamic.rsdiffbeforeafterboth--- a/crates/jrsonnet-evaluator/src/dynamic.rs
+++ b/crates/jrsonnet-evaluator/src/dynamic.rs
@@ -40,8 +40,9 @@
impl<T: Trace + Clone> ThunkValue for Pending<T> {
type Output = T;
- fn get(self: Box<Self>) -> Result<Self::Output> {
+ fn get(&self) -> Result<Self::Output> {
let Some(value) = self.0.get() else {
+ // TODO: Other error?
bail!(InfiniteRecursionDetected);
};
Ok(value.clone())
crates/jrsonnet-evaluator/src/obj.rsdiffbeforeafterboth--- a/crates/jrsonnet-evaluator/src/obj.rs
+++ b/crates/jrsonnet-evaluator/src/obj.rs
@@ -23,7 +23,7 @@
gc::WithCapacityExt as _,
identity_hash, in_frame,
operator::evaluate_add_op,
- val::ArrValue,
+ val::{ArrValue, ThunkValue},
CcUnbound, MaybeUnbound, Result, Thunk, Unbound, Val,
};
@@ -753,13 +753,45 @@
if !self.has_field_ex(key.clone(), true) {
return None;
}
- let obj = self.clone();
+ #[derive(Trace)]
+ struct ObjFieldThunk {
+ obj: ObjValue,
+ key: IStr,
+ }
+ impl ThunkValue for ObjFieldThunk {
+ type Output = Val;
- Some(Thunk!(move || Ok(obj.get(key)?.expect("field exists"))))
+ fn get(&self) -> Result<Self::Output> {
+ self.obj
+ .get(self.key.clone())
+ .transpose()
+ .expect("field existence checked")
+ }
+ }
+
+ Some(Thunk::new(ObjFieldThunk {
+ obj: self.clone(),
+ key,
+ }))
}
pub fn get_lazy_or_bail(&self, key: IStr) -> Thunk<Val> {
- let obj = self.clone();
- Thunk!(move || obj.get_or_bail(key))
+ #[derive(Trace)]
+ struct ObjFieldThunk {
+ obj: ObjValue,
+ key: IStr,
+ }
+ impl ThunkValue for ObjFieldThunk {
+ type Output = Val;
+
+ fn get(&self) -> Result<Self::Output> {
+ self.obj.get_or_bail(self.key.clone())
+ }
+ }
+
+ Thunk::new(ObjFieldThunk {
+ obj: self.clone(),
+ key,
+ })
}
pub fn ptr_eq(a: &Self, b: &Self) -> bool {
Cc::ptr_eq(&a.0, &b.0)
crates/jrsonnet-evaluator/src/val.rsdiffbeforeafterboth--- a/crates/jrsonnet-evaluator/src/val.rs
+++ b/crates/jrsonnet-evaluator/src/val.rs
@@ -2,13 +2,14 @@
cell::RefCell,
cmp::Ordering,
fmt::{self, Debug, Display},
+ marker::PhantomData,
mem::replace,
num::NonZeroU32,
ops::Deref,
rc::Rc,
};
-use jrsonnet_gcmodule::{Acyclic, Cc, Trace, TraceBox};
+use jrsonnet_gcmodule::{cc_dyn, Acyclic, Cc, Trace};
use jrsonnet_interner::IStr;
pub use jrsonnet_macros::Thunk;
use jrsonnet_types::ValType;
@@ -28,56 +29,106 @@
pub trait ThunkValue: Trace {
type Output;
- fn get(self: Box<Self>) -> Result<Self::Output>;
+ fn get(&self) -> Result<Self::Output>;
}
#[derive(Trace)]
-pub struct ThunkValueClosure<D: Trace, O: 'static> {
- env: D,
- // Carries no data, as it is not a real closure, all the
- // captured environment is stored in `env` field.
- #[trace(skip)]
- closure: fn(D) -> Result<O>,
+enum MemoizedClusureThunkInner<D: Trace, T: Trace> {
+ Computed(T),
+ Errored(Error),
+ Waiting {
+ env: D,
+ // Carries no data, as it is not a real closure, all the
+ // captured environment is stored in `env` field.
+ #[trace(skip)]
+ closure: fn(D) -> Result<T>,
+ },
+ Pending,
}
-impl<D: Trace, O: 'static> ThunkValueClosure<D, O> {
- pub fn new(env: D, closure: fn(D) -> Result<O>) -> Self {
- Self { env, closure }
+#[derive(Trace)]
+pub struct MemoizedClosureThunk<D: Trace, T: Trace>(RefCell<MemoizedClusureThunkInner<D, T>>);
+impl<D: Trace, T: Trace> MemoizedClosureThunk<D, T> {
+ pub fn new(env: D, closure: fn(D) -> Result<T>) -> Self {
+ Self(RefCell::new(MemoizedClusureThunkInner::Waiting {
+ env,
+ closure,
+ }))
}
}
-impl<D: Trace, O: 'static> ThunkValue for ThunkValueClosure<D, O> {
- type Output = O;
- fn get(self: Box<Self>) -> Result<Self::Output> {
- (self.closure)(self.env)
- }
-}
+impl<D: Trace, T: Trace + Clone> ThunkValue for MemoizedClosureThunk<D, T> {
+ type Output = T;
-#[derive(Trace)]
-enum ThunkInner<T: Trace> {
- Computed(T),
- Errored(Error),
- Waiting(TraceBox<dyn ThunkValue<Output = T>>),
- Pending,
+ fn get(&self) -> Result<Self::Output> {
+ match &*self.0.borrow() {
+ MemoizedClusureThunkInner::Computed(v) => return Ok(v.clone()),
+ MemoizedClusureThunkInner::Errored(e) => return Err(e.clone()),
+ MemoizedClusureThunkInner::Pending => return Err(InfiniteRecursionDetected.into()),
+ MemoizedClusureThunkInner::Waiting { .. } => (),
+ };
+ let MemoizedClusureThunkInner::Waiting { env, closure } = replace(
+ &mut *self.0.borrow_mut(),
+ MemoizedClusureThunkInner::Pending,
+ ) else {
+ unreachable!();
+ };
+ let new_value = match closure(env) {
+ Ok(v) => v,
+ Err(e) => {
+ *self.0.borrow_mut() = MemoizedClusureThunkInner::Errored(e.clone());
+ return Err(e);
+ }
+ };
+ *self.0.borrow_mut() = MemoizedClusureThunkInner::Computed(new_value.clone());
+ Ok(new_value)
+ }
}
-/// Lazily evaluated value
-#[allow(clippy::module_name_repetitions)]
-#[derive(Clone, Trace)]
-pub struct Thunk<T: Trace>(Cc<RefCell<ThunkInner<T>>>);
+cc_dyn!(
+ /// Lazily evaluated value
+ #[derive(Clone)] Thunk<V: Trace>,
+ ThunkValue<Output = V>,
+ pub fn new() {...}
+);
impl<T: Trace> Thunk<T> {
- pub fn evaluated(val: T) -> Self {
- Self(Cc::new(RefCell::new(ThunkInner::Computed(val))))
- }
- pub fn new(f: impl ThunkValue<Output = T> + 'static) -> Self {
- Self(Cc::new(RefCell::new(ThunkInner::Waiting(TraceBox(
- Box::new(f),
- )))))
+ pub fn evaluated(val: T) -> Self
+ where
+ T: Clone,
+ {
+ #[derive(Trace)]
+ struct EvaluatedThunk<T: Trace>(T);
+ impl<T> ThunkValue for EvaluatedThunk<T>
+ where
+ T: Clone + Trace,
+ {
+ type Output = T;
+
+ fn get(&self) -> Result<Self::Output> {
+ Ok(self.0.clone())
+ }
+ }
+ Self::new(EvaluatedThunk(val))
}
pub fn errored(e: Error) -> Self {
- Self(Cc::new(RefCell::new(ThunkInner::Errored(e))))
+ #[derive(Trace)]
+ struct ErroredThunk<T: Trace>(Error, PhantomData<T>);
+ impl<T> ThunkValue for ErroredThunk<T>
+ where
+ T: Trace,
+ {
+ type Output = T;
+
+ fn get(&self) -> Result<Self::Output> {
+ Err(self.0.clone())
+ }
+ }
+ Self::new(ErroredThunk(e, PhantomData))
}
- pub fn result(res: Result<T, Error>) -> Self {
+ pub fn result(res: Result<T, Error>) -> Self
+ where
+ T: Clone,
+ {
match res {
Ok(o) => Self::evaluated(o),
Err(e) => Self::errored(e),
@@ -101,25 +152,7 @@
/// - Lazy value evaluation returned error
/// - This method was called during inner value evaluation
pub fn evaluate(&self) -> Result<T> {
- match &*self.0.borrow() {
- ThunkInner::Computed(v) => return Ok(v.clone()),
- ThunkInner::Errored(e) => return Err(e.clone()),
- ThunkInner::Pending => return Err(InfiniteRecursionDetected.into()),
- ThunkInner::Waiting(..) => (),
- };
- let ThunkInner::Waiting(value) = replace(&mut *self.0.borrow_mut(), ThunkInner::Pending)
- else {
- unreachable!();
- };
- let new_value = match value.0.get() {
- Ok(v) => v,
- Err(e) => {
- *self.0.borrow_mut() = ThunkInner::Errored(e.clone());
- return Err(e);
- }
- };
- *self.0.borrow_mut() = ThunkInner::Computed(new_value.clone());
- Ok(new_value)
+ self.0.get()
}
}
@@ -134,7 +167,7 @@
pub fn map<M>(self, mapper: M) -> Thunk<M::Output>
where
M: ThunkMapper<Input>,
- M::Output: Trace,
+ M::Output: Trace + Clone,
{
let inner = self;
Thunk!(move || {
@@ -145,7 +178,7 @@
}
}
-impl<T: Trace> From<Result<T>> for Thunk<T> {
+impl<T: Trace + Clone> From<Result<T>> for Thunk<T> {
fn from(value: Result<T>) -> Self {
match value {
Ok(o) => Self::evaluated(o),
@@ -162,7 +195,7 @@
}
}
-impl<T: Trace + Default> Default for Thunk<T> {
+impl<T: Trace + Default + Clone> Default for Thunk<T> {
fn default() -> Self {
Self::evaluated(T::default())
}
crates/jrsonnet-formatter/src/comments.rsdiffbeforeafterboth--- a/crates/jrsonnet-formatter/src/comments.rs
+++ b/crates/jrsonnet-formatter/src/comments.rs
@@ -73,7 +73,10 @@
p!(out, str(" "));
}
p!(out, str("/* ") string(lines[0].trim().to_string()) str(" */"));
- if matches!(loc, CommentLocation::AboveItem | CommentLocation::EndOfItems) {
+ if matches!(
+ loc,
+ CommentLocation::AboveItem | CommentLocation::EndOfItems
+ ) {
p!(out, nl);
}
} else if !lines.is_empty() {
crates/jrsonnet-macros/src/lib.rsdiffbeforeafterboth--- a/crates/jrsonnet-macros/src/lib.rs
+++ b/crates/jrsonnet-macros/src/lib.rs
@@ -879,6 +879,6 @@
quote! {{
#move_check
#(#trace_check)*
- ::jrsonnet_evaluator::Thunk::new(::jrsonnet_evaluator::val::ThunkValueClosure::new(#env, #closure))
+ ::jrsonnet_evaluator::Thunk::new(::jrsonnet_evaluator::val::MemoizedClosureThunk::new(#env, #closure))
}}.into()
}
resultdiffbeforeafterboth--- a/result
+++ /dev/null
@@ -1 +0,0 @@
-/nix/store/nd6v7jksg1dqhpx4x4vqgy5ry1nkb9lk-jrsonnet-current
\ No newline at end of file
xtask/src/sourcegen/mod.rsdiffbeforeafterboth1use std::{collections::HashMap, path::PathBuf};23use anyhow::Result;4use ast::{lower, AstSrc};5use itertools::Itertools;6use kinds::{KindsSrc, TokenKind};7use proc_macro2::{Ident, Punct, Spacing, Span, TokenStream};8use quote::{format_ident, quote};9use ungrammar::Grammar;10use util::{ensure_file_contents, reformat, to_pascal_case, to_upper_snake_case};1112mod ast;13mod kinds;14mod util;1516enum SpecialName {17 Literal,18 Meta,19 Error,20}21fn classify_special(name: &str) -> Option<(SpecialName, &str)> {22 let name = name.strip_suffix('!')?;23 Some(if let Some(name) = name.strip_prefix("LIT_") {24 (SpecialName::Literal, name)25 } else if let Some(name) = name.strip_prefix("META_") {26 (SpecialName::Meta, name)27 } else if let Some(name) = name.strip_prefix("ERROR_") {28 (SpecialName::Error, name)29 } else {30 return None;31 })32}3334pub fn generate_ungrammar() -> Result<()> {35 let grammar: Grammar = include_str!(concat!(36 env!("CARGO_MANIFEST_DIR"),37 "/../crates/jrsonnet-rowan-parser/jsonnet.ungram"38 ))39 .parse()?;4041 let mut kinds = kinds::jsonnet_kinds();42 let ast = lower(&kinds, &grammar);4344 for token in grammar.tokens() {45 let token = &grammar[token];46 let token = &token.name.clone();47 if !kinds.is_token(token) {48 if let Some((special, name)) = classify_special(token) {49 match special {50 SpecialName::Literal => panic!("literal is not defined: {name}"),51 SpecialName::Meta => {52 eprintln!("implicit meta: {name}");53 kinds.define_token(TokenKind::Meta {54 grammar_name: token.to_owned(),55 name: format!("META_{name}"),56 });57 }58 SpecialName::Error => {59 eprintln!("implicit error: {name}");60 kinds.define_token(TokenKind::Error {61 grammar_name: token.to_owned(),62 name: format!("ERROR_{name}"),63 regex: None,64 priority: None,65 is_lexer_error: true,66 });67 }68 };69 continue;70 };71 let name = to_upper_snake_case(token);72 eprintln!("implicit kw: {token}");73 kinds.define_token(TokenKind::Keyword {74 code: token.to_owned(),75 name: format!("{name}_KW"),76 });77 }78 }79 for node in &ast.nodes {80 let name = to_upper_snake_case(&node.name);81 kinds.define_node(&name);82 }83 for enum_ in &ast.enums {84 let name = to_upper_snake_case(&enum_.name);85 kinds.define_node(&name);86 }87 for token_enum in &ast.token_enums {88 let name = to_upper_snake_case(&token_enum.name);89 kinds.define_node(&name);90 }9192 let syntax_kinds = generate_syntax_kinds(&kinds, &ast)?;9394 let nodes = generate_nodes(&kinds, &ast)?;95 ensure_file_contents(96 &PathBuf::from(concat!(97 env!("CARGO_MANIFEST_DIR"),98 "/../crates/jrsonnet-rowan-parser/src/generated/syntax_kinds.rs",99 )),100 &syntax_kinds,101 );102 ensure_file_contents(103 &PathBuf::from(concat!(104 env!("CARGO_MANIFEST_DIR"),105 "/../crates/jrsonnet-rowan-parser/src/generated/nodes.rs",106 )),107 &nodes,108 );109 Ok(())110}111112fn generate_syntax_kinds(kinds: &KindsSrc, grammar: &AstSrc) -> Result<String> {113 let t_macros = kinds.tokens().filter_map(TokenKind::expand_t_macros);114 let token_kinds = kinds.tokens().map(TokenKind::expand_kind);115116 let keywords = kinds117 .tokens()118 .filter(|k| matches!(k, TokenKind::Keyword { .. }))119 .map(TokenKind::name)120 .map(|n| format_ident!("{n}"));121122 let nodes = kinds123 .nodes124 .iter()125 .map(|name| format_ident!("{}", name))126 .collect::<Vec<_>>();127128 let enums = grammar129 .enums130 .iter()131 .map(|e| format_ident!("{}", to_upper_snake_case(&e.name)))132 .chain(133 grammar134 .token_enums135 .iter()136 .map(|e| format_ident!("{}", to_upper_snake_case(&e.name))),137 );138139 let ast = quote! {140 #![allow(bad_style, missing_docs, unreachable_pub, clippy::manual_non_exhaustive, clippy::match_like_matches_macro)]141 use logos::Logos;142143 /// The kind of syntax node, e.g. `IDENT`, `USE_KW`, or `STRUCT`.144 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Logos)]145 #[repr(u16)]146 pub enum SyntaxKind {147 #[doc(hidden)]148 TOMBSTONE,149 #[doc(hidden)]150 EOF,151 #(#token_kinds,)*152 LEXING_ERROR,153 __LAST_TOKEN,154 #(#nodes,)*155 #[doc(hidden)]156 __LAST,157 }158 use self::SyntaxKind::*;159160 impl SyntaxKind {161 pub fn is_keyword(self) -> bool {162 match self {163 #(#keywords)|* => true,164 _ => false,165 }166 }167 pub fn is_enum(self) -> bool {168 match self {169 #(#enums)|* => true,170 _ => false,171 }172 }173174 pub fn from_raw(r: u16) -> Self {175 assert!(r < Self::__LAST as u16);176 unsafe { std::mem::transmute(r) }177 }178 pub fn into_raw(self) -> u16 {179 self as u16180 }181 }182183 #[macro_export]184 macro_rules! T {#(#t_macros);*}185 #[allow(unused_imports)]186 pub use T;187 };188189 reformat(&ast.to_string())190}191192#[allow(clippy::too_many_lines, clippy::cognitive_complexity)]193fn generate_nodes(kinds: &KindsSrc, grammar: &AstSrc) -> Result<String> {194 let (node_defs, node_boilerplate_impls): (Vec<_>, Vec<_>) = grammar195 .nodes196 .iter()197 .map(|node| {198 let name = format_ident!("{}", node.name);199 let kind = format_ident!("{}", to_upper_snake_case(&node.name));200 let traits = node.traits.iter().map(|trait_name| {201 let trait_name = format_ident!("{}", trait_name);202 quote!(impl ast::#trait_name for #name {})203 });204205 let mut type_positions: HashMap<String, usize> = HashMap::new();206 let field_positions: Vec<_> = node.fields.iter().map(|field| {207 let ty_str = field.ty().to_string();208 let pos = *type_positions.get(&ty_str).unwrap_or(&0);209 type_positions.insert(ty_str, pos + 1);210 pos211 }).collect();212213 let methods = node.fields.iter().zip(field_positions.iter()).map(|(field, &pos)| {214 let method_name = field.method_name(kinds);215 let ty = field.ty();216217 if field.is_many() {218 quote! {219 pub fn #method_name(&self) -> AstChildren<#ty> {220 support::children(&self.syntax)221 }222 }223 } else if let Some(token_kind) = field.token_kind(kinds) {224 quote! {225 pub fn #method_name(&self) -> Option<#ty> {226 support::token(&self.syntax, #token_kind)227 }228 }229 } else if field.is_token_enum(grammar) {230 quote! {231 pub fn #method_name(&self) -> Option<#ty> {232 support::token_child(&self.syntax)233 }234 }235 } else if pos == 0 {236 quote! {237 pub fn #method_name(&self) -> Option<#ty> {238 support::children(&self.syntax).next()239 }240 }241 } else {242 quote! {243 pub fn #method_name(&self) -> Option<#ty> {244 support::children(&self.syntax).nth(#pos)245 }246 }247 }248 });249 (250 quote! {251 #[pretty_doc_comment_placeholder_workaround]252 #[derive(Debug, Clone, PartialEq, Eq, Hash)]253 pub struct #name {254 pub(crate) syntax: SyntaxNode,255 }256257 #(#traits)*258259 impl #name {260 #(#methods)*261 }262 },263 quote! {264 impl AstNode for #name {265 fn can_cast(kind: SyntaxKind) -> bool {266 kind == #kind267 }268 fn cast(syntax: SyntaxNode) -> Option<Self> {269 if Self::can_cast(syntax.kind()) { Some(Self { syntax }) } else { None }270 }271 fn syntax(&self) -> &SyntaxNode { &self.syntax }272 }273 },274 )275 })276 .unzip();277278 let (enum_defs, enum_boilerplate_impls): (Vec<_>, Vec<_>) = grammar279 .enums280 .iter()281 .map(|en| {282 let variants: Vec<_> = en283 .variants284 .iter()285 .map(|var| format_ident!("{}", var))286 .collect();287 let name = format_ident!("{}", en.name);288 let kinds: Vec<_> = variants289 .iter()290 .map(|name| format_ident!("{}", to_upper_snake_case(&name.to_string())))291 .collect();292 let traits = en.traits.iter().map(|trait_name| {293 let trait_name = format_ident!("{}", trait_name);294 quote!(impl ast::#trait_name for #name {})295 });296297 let ast_node = quote! {298 impl AstNode for #name {299 fn can_cast(kind: SyntaxKind) -> bool {300 match kind {301 #(#kinds)|* => true,302 _ => false,303 }304 }305 fn cast(syntax: SyntaxNode) -> Option<Self> {306 let res = match syntax.kind() {307 #(308 #kinds => #name::#variants(#variants { syntax }),309 )*310 _ => return None,311 };312 Some(res)313 }314 fn syntax(&self) -> &SyntaxNode {315 match self {316 #(317 #name::#variants(it) => &it.syntax,318 )*319 }320 }321 }322 };323324 (325 quote! {326 #[pretty_doc_comment_placeholder_workaround]327 #[derive(Debug, Clone, PartialEq, Eq, Hash)]328 pub enum #name {329 #(#variants(#variants),)*330 }331332 #(#traits)*333 },334 quote! {335 #(336 impl From<#variants> for #name {337 fn from(node: #variants) -> #name {338 #name::#variants(node)339 }340 }341 )*342 #ast_node343 },344 )345 })346 .unzip();347348 let (token_enum_defs, token_enum_boilerplate_impls): (Vec<_>, Vec<_>) = grammar349 .token_enums350 .iter()351 .map(|en| {352 let variants: Vec<_> = en353 .variants354 .iter()355 .map(|token| {356 format_ident!(357 "{}",358 to_pascal_case(kinds.token(token).expect("token exists").name())359 )360 })361 .collect();362 let name = format_ident!("{}", en.name);363 let kind_name = format_ident!("{}Kind", en.name);364 let kinds: Vec<_> = variants365 .iter()366 .map(|name| format_ident!("{}", to_upper_snake_case(&name.to_string())))367 .collect();368369 let ast_node = quote! {370 impl AstToken for #name {371 fn can_cast(kind: SyntaxKind) -> bool {372 #kind_name::can_cast(kind)373 }374 fn cast(syntax: SyntaxToken) -> Option<Self> {375 let kind = #kind_name::cast(syntax.kind())?;376 Some(#name { syntax, kind })377 }378 fn syntax(&self) -> &SyntaxToken {379 &self.syntax380 }381 }382383 impl #kind_name {384 fn can_cast(kind: SyntaxKind) -> bool {385 match kind {386 #(#kinds)|* => true,387 _ => false,388 }389 }390 pub fn cast(kind: SyntaxKind) -> Option<Self> {391 let res = match kind {392 #(#kinds => Self::#variants,)*393 _ => return None,394 };395 Some(res)396 }397 }398 };399400 (401 quote! {402 #[pretty_doc_comment_placeholder_workaround]403 #[derive(Debug, Clone, PartialEq, Eq, Hash)]404 pub struct #name { syntax: SyntaxToken, kind: #kind_name }405406 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]407 pub enum #kind_name {408 #(#variants,)*409 }410 },411 quote! {412 #ast_node413414 impl #name {415 pub fn kind(&self) -> #kind_name {416 self.kind417 }418 }419420 impl std::fmt::Display for #name {421 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {422 std::fmt::Display::fmt(self.syntax(), f)423 }424 }425 },426 )427 })428 .unzip();429430 let (any_node_defs, any_node_boilerplate_impls): (Vec<_>, Vec<_>) = grammar431 .nodes432 .iter()433 .flat_map(|node| node.traits.iter().map(move |t| (t, node)))434 .into_group_map()435 .into_iter()436 .sorted_by_key(|(k, _)| *k)437 .map(|(trait_name, nodes)| {438 let name = format_ident!("Any{}", trait_name);439 let trait_name = format_ident!("{}", trait_name);440 let kinds: Vec<_> = nodes441 .iter()442 .map(|name| format_ident!("{}", to_upper_snake_case(&name.name.to_string())))443 .collect();444445 (446 quote! {447 #[pretty_doc_comment_placeholder_workaround]448 #[derive(Debug, Clone, PartialEq, Eq, Hash)]449 pub struct #name {450 pub(crate) syntax: SyntaxNode,451 }452 impl ast::#trait_name for #name {}453 },454 quote! {455 impl #name {456 #[inline]457 pub fn new<T: ast::#trait_name>(node: T) -> #name {458 #name {459 syntax: node.syntax().clone()460 }461 }462 }463 impl AstNode for #name {464 fn can_cast(kind: SyntaxKind) -> bool {465 match kind {466 #(#kinds)|* => true,467 _ => false,468 }469 }470 fn cast(syntax: SyntaxNode) -> Option<Self> {471 Self::can_cast(syntax.kind()).then(|| #name { syntax })472 }473 fn syntax(&self) -> &SyntaxNode {474 &self.syntax475 }476 }477 },478 )479 })480 .unzip();481482 let enum_names = grammar.enums.iter().map(|it| &it.name);483 let node_names = grammar.nodes.iter().map(|it| &it.name);484485 let display_impls = enum_names486 .chain(node_names.clone())487 .map(|it| format_ident!("{}", it))488 .map(|name| {489 quote! {490 impl std::fmt::Display for #name {491 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {492 std::fmt::Display::fmt(self.syntax(), f)493 }494 }495 }496 });497498 let ast = quote! {499 #![allow(non_snake_case, clippy::match_like_matches_macro)]500501 use crate::{502 SyntaxNode, SyntaxToken, SyntaxKind::{self, *},503 ast::{AstNode, AstToken, AstChildren, support},504 T,505 };506507 #(#node_defs)*508 #(#enum_defs)*509 #(#token_enum_defs)*510 #(#any_node_defs)*511 #(#node_boilerplate_impls)*512 #(#enum_boilerplate_impls)*513 #(#token_enum_boilerplate_impls)*514 #(#any_node_boilerplate_impls)*515 #(#display_impls)*516 };517518 let ast = ast.to_string().replace("T ! [", "T![");519520 let mut res = String::with_capacity(ast.len() * 2);521522 let mut docs = grammar523 .nodes524 .iter()525 .map(|it| &it.doc)526 .chain(grammar.enums.iter().map(|it| &it.doc));527528 for chunk in ast.split("# [pretty_doc_comment_placeholder_workaround] ") {529 res.push_str(chunk);530 if let Some(doc) = docs.next() {531 write_doc_comment(doc, &mut res);532 }533 }534535 let res = reformat(&res)?;536 Ok(res.replace("#[derive", "\n#[derive"))537}538539fn write_doc_comment(contents: &[String], dest: &mut String) {540 use std::fmt::Write;541 for line in contents {542 writeln!(dest, "///{line}").unwrap();543 }544}545546pub fn escape_token_macro(token: &str) -> TokenStream {547 if "{}[]()$".contains(token) {548 let c = token.chars().next().unwrap();549 quote! { #c }550 } else if token.contains(|v| v == '$') {551 quote! { #token }552 } else if token.chars().all(|v| ('a'..='z').contains(&v)) {553 let i = Ident::new(&token, Span::call_site());554 quote! { #i }555 } else {556 let cs = token.chars().map(|c| Punct::new(c, Spacing::Joint));557 quote! { #(#cs)* }558 }559}