1use itertools::Itertools;
2use spade_common::{
3 location_info::Loc,
4 name::{Identifier, NameID},
5};
6use spade_types::meta_types::MetaType;
7
8use crate::{
9 symbol_table::GenericArg, ConstGeneric, Parameter, ParameterList, TraitName, TraitSpec,
10 TypeExpression, TypeParam, TypeSpec, UnitHead, UnitKind,
11};
12
13pub trait MaybePrettyPrint {
14 fn maybe_pretty_print(&self) -> Option<String>;
15
16 fn with_trailing_space(&self) -> String {
17 match self.maybe_pretty_print() {
18 Some(s) => format!("{} ", s),
19 None => "".to_string(),
20 }
21 }
22}
23
24pub trait PrettyPrint {
25 fn pretty_print(&self) -> String;
26}
27
28impl PrettyPrint for NameID {
29 fn pretty_print(&self) -> String {
30 format!("{}", self.1.tail())
31 }
32}
33
34impl PrettyPrint for Identifier {
35 fn pretty_print(&self) -> String {
36 format!("{self}")
37 }
38}
39
40impl MaybePrettyPrint for MetaType {
41 fn maybe_pretty_print(&self) -> Option<String> {
42 match self {
43 MetaType::Any => Some("#any".to_string()),
44 MetaType::Type => None,
45 MetaType::Number => Some("#number".to_string()),
46 MetaType::Int => Some("#int".to_string()),
47 MetaType::Uint => Some("#uint".to_string()),
48 MetaType::Bool => Some("#bool".to_string()),
49 MetaType::Str => Some("#str".to_string()),
50 }
51 }
52}
53
54impl PrettyPrint for GenericArg {
55 fn pretty_print(&self) -> String {
56 match self {
57 GenericArg::TypeName { name, traits } => {
58 let traits = if traits.is_empty() {
59 "".to_string()
60 } else {
61 format!(
62 ": {}",
63 traits
64 .iter()
65 .map(|t| format!("{}", t.pretty_print()))
66 .join(", ")
67 )
68 };
69
70 format!("{}{}", name, traits)
71 }
72 GenericArg::TypeWithMeta { name, meta } => {
73 let meta = match meta {
74 spade_types::meta_types::MetaType::Any => "#any ",
75 spade_types::meta_types::MetaType::Type => "",
76 spade_types::meta_types::MetaType::Number => "#number ",
77 spade_types::meta_types::MetaType::Int => "#int ",
78 spade_types::meta_types::MetaType::Uint => "#uint ",
79 spade_types::meta_types::MetaType::Bool => "#bool ",
80 spade_types::meta_types::MetaType::Str => "#str ",
81 };
82
83 format!("{}{}", meta, name.pretty_print())
84 }
85 }
86 }
87}
88
89impl PrettyPrint for TraitName {
90 fn pretty_print(&self) -> String {
91 match self {
92 TraitName::Named(name) => name.pretty_print(),
93 TraitName::Anonymous(_) => "[Anonymous]".to_string(),
94 }
95 }
96}
97
98impl PrettyPrint for TraitSpec {
99 fn pretty_print(&self) -> String {
100 let tp = match &self.type_params {
101 Some(tp) => format!("<{}>", tp.iter().map(|tp| tp.pretty_print()).join(", ")),
102 None => "".to_string(),
103 };
104 format!("{}{}", self.name.pretty_print(), tp)
105 }
106}
107impl PrettyPrint for ConstGeneric {
108 fn pretty_print(&self) -> String {
109 match self {
110 ConstGeneric::Name(n) => n.pretty_print(),
111 ConstGeneric::Int(big_int) => format!("{big_int}"),
112 ConstGeneric::Str(s) => format!("{s:?}"),
113 ConstGeneric::Add(lhs, rhs) => {
114 format!("({} + {})", lhs.pretty_print(), rhs.pretty_print())
115 }
116 ConstGeneric::Sub(lhs, rhs) => {
117 format!("({} - {})", lhs.pretty_print(), rhs.pretty_print())
118 }
119 ConstGeneric::Mul(lhs, rhs) => {
120 format!("({} * {})", lhs.pretty_print(), rhs.pretty_print())
121 }
122 ConstGeneric::Div(lhs, rhs) => {
123 format!("({} / {})", lhs.pretty_print(), rhs.pretty_print())
124 }
125 ConstGeneric::Mod(lhs, rhs) => {
126 format!("({} % {})", lhs.pretty_print(), rhs.pretty_print())
127 }
128 ConstGeneric::UintBitsToFit(inner) => {
129 format!("{}", inner.pretty_print())
130 }
131 ConstGeneric::Eq(lhs, rhs) => {
132 format!("({} == {})", lhs.pretty_print(), rhs.pretty_print())
133 }
134 ConstGeneric::NotEq(lhs, rhs) => {
135 format!("({} != {})", lhs.pretty_print(), rhs.pretty_print())
136 }
137 }
138 }
139}
140
141impl PrettyPrint for TypeExpression {
142 fn pretty_print(&self) -> String {
143 match self {
144 TypeExpression::Integer(val) => format!("{val}"),
145 TypeExpression::String(val) => format!("{val:?}"),
146 TypeExpression::TypeSpec(ts) => ts.pretty_print(),
147 TypeExpression::ConstGeneric(cg) => cg.pretty_print(),
148 }
149 }
150}
151
152impl PrettyPrint for TypeSpec {
153 fn pretty_print(&self) -> String {
154 match self {
155 TypeSpec::Declared(base, args) => {
156 let args = if !args.is_empty() {
157 format!("<{}>", args.iter().map(|arg| arg.pretty_print()).join(", "))
158 } else {
159 "".to_string()
160 };
161 format!("{}{}", base.pretty_print(), args)
162 }
163 TypeSpec::Generic(name) => name.pretty_print(),
164 TypeSpec::Tuple(inner) => format!(
165 "({})",
166 inner.iter().map(|arg| arg.pretty_print()).join(", ")
167 ),
168 TypeSpec::Array { inner, size } => {
169 format!("[{}; {}]", inner.pretty_print(), size.pretty_print())
170 }
171 TypeSpec::Inverted(inner) => format!("inv {}", inner.pretty_print()),
172 TypeSpec::Wire(inner) => format!("&{}", inner.pretty_print()),
173 TypeSpec::TraitSelf(_) => format!("self"),
174 TypeSpec::Wildcard(_) => format!("_"),
175 }
176 }
177}
178
179impl PrettyPrint for TypeParam {
180 fn pretty_print(&self) -> String {
181 let Self {
182 ident: _,
183 name_id,
184 trait_bounds,
185 meta,
186 } = self;
187
188 let traits = if trait_bounds.is_empty() {
189 "".to_string()
190 } else {
191 format!(
192 "<{}>",
193 trait_bounds.iter().map(|tb| tb.pretty_print()).join(", ")
194 )
195 };
196
197 format!(
198 "{}{}{}",
199 meta.with_trailing_space(),
200 name_id.pretty_print(),
201 traits
202 )
203 }
204}
205
206impl PrettyPrint for UnitKind {
207 fn pretty_print(&self) -> String {
208 match self {
209 UnitKind::Function(crate::FunctionKind::Fn) => "fn".to_string(),
210 UnitKind::Function(crate::FunctionKind::Struct) => "struct".to_string(),
211 UnitKind::Function(crate::FunctionKind::Enum) => "enum variant".to_string(),
212 UnitKind::Entity => "entity".to_string(),
213 UnitKind::Pipeline {
214 depth,
215 depth_typeexpr_id: _,
216 } => format!("pipeline({})", depth.pretty_print()),
217 }
218 }
219}
220
221impl PrettyPrint for UnitHead {
222 fn pretty_print(&self) -> String {
223 let Self {
224 name,
225 inputs,
226 is_nonstatic_method: _,
227 output_type,
228 unit_type_params,
229 scope_type_params: _,
230 unit_kind,
231 where_clauses: _,
232 unsafe_marker,
233 documentation: _,
234 } = self;
235 let output_type = match output_type {
236 Some(output_type) => format!(" -> {}", output_type.pretty_print()),
237 None => "".to_string(),
238 };
239 let type_params = if unit_type_params.is_empty() {
240 "".to_string()
241 } else {
242 format!(
243 "<{}>",
244 unit_type_params
245 .iter()
246 .map(|tp| tp.pretty_print())
247 .join(", ")
248 )
249 };
250 let inputs = inputs.pretty_print();
251 format!(
252 "{}{} {}{}({}){}",
253 if unsafe_marker.is_some() {
254 "unsafe "
255 } else {
256 ""
257 },
258 unit_kind.pretty_print(),
259 name,
260 type_params,
261 inputs,
262 output_type
263 )
264 }
265}
266
267impl PrettyPrint for Parameter {
268 fn pretty_print(&self) -> String {
269 let Parameter {
270 no_mangle: _,
271 field_translator: _,
272 name,
273 ty,
274 } = self;
275
276 format!("{}: {}", name.pretty_print(), ty.pretty_print())
277 }
278}
279
280impl PrettyPrint for ParameterList {
281 fn pretty_print(&self) -> String {
282 self.0.iter().map(|param| param.pretty_print()).join(", ")
283 }
284}
285
286impl<T> PrettyPrint for &T
287where
288 T: PrettyPrint,
289{
290 fn pretty_print(&self) -> String {
291 (*self).pretty_print()
292 }
293}
294
295impl<T> PrettyPrint for Loc<T>
296where
297 T: PrettyPrint,
298{
299 fn pretty_print(&self) -> String {
300 self.inner.pretty_print()
301 }
302}
303
304impl<T> PrettyPrint for Option<T>
305where
306 T: PrettyPrint,
307{
308 fn pretty_print(&self) -> String {
309 match self {
310 Some(inner) => inner.pretty_print(),
311 None => String::new(),
312 }
313 }
314}