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