power-instruction-analyzer-proc-macro/src/lib.rs

   1 // SPDX-License-Identifier: LGPL-2.1-or-later
   2 // See Notices.txt for copyright information
   3
   4 use proc_macro2::{Ident, Span, TokenStream};
   5 use quote::{quote, ToTokens, TokenStreamExt};
   6 use std::{borrow::Cow, fmt, fmt::Write};
   7 use syn::{
   8     braced, bracketed, parenthesized,
   9     parse::{Parse, ParseStream},
  10     parse_macro_input,
  11     punctuated::Punctuated,
  12     Attribute, Error, ItemFn, LitStr, Token,
  13 };
  14
  15 macro_rules! valid_enumerants_as_string {
  16     ($enumerant:ident) => {
  17         concat!("`", stringify!($enumerant), "`")
  18     };
  19     ($enumerant1:ident, $enumerant2:ident) => {
  20         concat!("`", stringify!($enumerant1), "` and `", stringify!($enumerant2), "`")
  21     };
  22     ($($enumerant:ident),+) => {
  23         valid_enumerants_as_string!((), ($($enumerant),+))
  24     };
  25     (($first_enumerant:ident, $($enumerant:ident,)+), ($last_enumerant:ident)) => {
  26         concat!(
  27             "`",
  28             stringify!($first_enumerant),
  29             $(
  30                 "`, `",
  31                 stringify!($enumerant),
  32             )+
  33             "`, and `",
  34             stringify!($last_enumerant),
  35             "`"
  36         )
  37     };
  38     (($($enumerants:ident,)*), ($next_enumerant:ident, $($rest:ident),*)) => {
  39         valid_enumerants_as_string!(($($enumerants,)* $next_enumerant,), ($($rest),*))
  40     };
  41     () => {
  42         "<nothing>"
  43     };
  44 }
  45
  46 macro_rules! ident_enum {
  47     (
  48         #[parse_error_msg = $parse_error_msg:literal]
  49         enum $enum_name:ident {
  50             $(
  51                 $enumerant:ident,
  52             )*
  53         }
  54     ) => {
  55         #[derive(Copy, Clone, Eq, PartialEq, Hash)]
  56         enum $enum_name<T = Span> {
  57             $(
  58                 $enumerant(T),
  59             )*
  60         }
  61
  62         impl<T> fmt::Debug for $enum_name<T> {
  63             fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  64                 f.write_str(self.name())
  65             }
  66         }
  67
  68         impl<T> $enum_name<T> {
  69             fn enumerant(&self) -> $enum_name<()> {
  70                 match self {
  71                     $(
  72                         $enum_name::$enumerant(_) => $enum_name::$enumerant(()),
  73                     )*
  74                 }
  75             }
  76             fn name(&self) -> &'static str {
  77                 match self {
  78                     $(
  79                         $enum_name::$enumerant(_) => stringify!($enumerant),
  80                     )*
  81                 }
  82             }
  83         }
  84
  85         impl $enum_name {
  86             fn into_ident(self) -> Ident {
  87                 match self {
  88                     $(
  89                         $enum_name::$enumerant(span) => Ident::new(stringify!($enumerant), span),
  90                     )*
  91                 }
  92             }
  93         }
  94
  95         impl ToTokens for $enum_name<Span> {
  96             fn to_tokens(&self, tokens: &mut TokenStream) {
  97                 tokens.append(self.clone().into_ident());
  98             }
  99         }
 100
 101         impl Parse for $enum_name<Span> {
 102             fn parse(input: ParseStream) -> syn::Result<Self> {
 103                 let id: Ident = input.parse()?;
 104                 $(
 105                     if id == stringify!($enumerant) {
 106                         return Ok($enum_name::$enumerant(id.span()));
 107                     }
 108                 )*
 109                 Err(Error::new_spanned(
 110                     id,
 111                     concat!(
 112                         $parse_error_msg,
 113                         ": valid values are: ",
 114                         valid_enumerants_as_string!($($enumerant),*)
 115                     )
 116                 ))
 117             }
 118         }
 119     };
 120 }
 121
 122 ident_enum! {
 123     #[parse_error_msg = "unknown instruction input"]
 124     enum InstructionInput {
 125         Ra,
 126         Rb,
 127         Rc,
 128         Carry,
 129     }
 130 }
 131
 132 ident_enum! {
 133     #[parse_error_msg = "unknown instruction output"]
 134     enum InstructionOutput {
 135         Rt,
 136         Carry,
 137         Overflow,
 138         CR0,
 139     }
 140 }
 141
 142 #[derive(Debug, Clone)]
 143 enum AssemblyTextFragment {
 144     Text(String),
 145     InputIndex(usize),
 146     OutputIndex(usize),
 147 }
 148
 149 struct InlineAssembly {
 150     text: Vec<AssemblyTextFragment>,
 151     text_span: Span,
 152     inputs: Vec<TokenStream>,
 153     outputs: Vec<TokenStream>,
 154     clobbers: Vec<TokenStream>,
 155 }
 156
 157 impl fmt::Write for InlineAssembly {
 158     fn write_str(&mut self, s: &str) -> fmt::Result {
 159         if let Some(AssemblyTextFragment::Text(v)) = self.text.last_mut() {
 160             *v += s;
 161         } else {
 162             self.text.push(AssemblyTextFragment::Text(String::from(s)));
 163         }
 164         Ok(())
 165     }
 166 }
 167
 168 impl InlineAssembly {
 169     fn new(text_span: Span) -> Self {
 170         Self {
 171             text: Vec::new(),
 172             text_span,
 173             inputs: Vec::new(),
 174             outputs: Vec::new(),
 175             clobbers: Vec::new(),
 176         }
 177     }
 178     fn to_text(&self) -> String {
 179         let mut retval = String::new();
 180         for text in &self.text {
 181             match text {
 182                 AssemblyTextFragment::Text(text) => retval += text,
 183                 AssemblyTextFragment::InputIndex(index) => {
 184                     write!(retval, "{}", index + self.outputs.len()).unwrap();
 185                 }
 186                 AssemblyTextFragment::OutputIndex(index) => write!(retval, "{}", index).unwrap(),
 187             }
 188         }
 189         retval
 190     }
 191     fn write_input_index(&mut self, index: usize) -> fmt::Result {
 192         self.text.push(AssemblyTextFragment::InputIndex(index));
 193         Ok(())
 194     }
 195     fn write_output_index(&mut self, index: usize) -> fmt::Result {
 196         self.text.push(AssemblyTextFragment::OutputIndex(index));
 197         Ok(())
 198     }
 199 }
 200
 201 impl ToTokens for InlineAssembly {
 202     fn to_tokens(&self, tokens: &mut TokenStream) {
 203         let Self {
 204             text: _,
 205             text_span,
 206             inputs,
 207             outputs,
 208             clobbers,
 209         } = self;
 210         let text = LitStr::new(&self.to_text(), text_span.clone());
 211         let value = quote! {
 212             llvm_asm!(#text : #(#outputs),* : #(#inputs),* : #(#clobbers),*)
 213         };
 214         value.to_tokens(tokens);
 215     }
 216 }
 217
 218 #[derive(Debug)]
 219 struct Instruction {
 220     enumerant: Ident,
 221     fn_name: Ident,
 222     inputs: Punctuated<InstructionInput, Token!(,)>,
 223     outputs: Punctuated<InstructionOutput, Token!(,)>,
 224     instruction_name: LitStr,
 225 }
 226
 227 impl Parse for Instruction {
 228     fn parse(input: ParseStream) -> syn::Result<Self> {
 229         input.parse::<Token!(#)>()?;
 230         let enumerant_attr_tokens;
 231         bracketed!(enumerant_attr_tokens in input);
 232         let enumerant_name: Ident = enumerant_attr_tokens.parse()?;
 233         if enumerant_name != "enumerant" {
 234             return Err(Error::new_spanned(
 235                 enumerant_name,
 236                 "expected `#[enumerant = ...]` attribute",
 237             ));
 238         }
 239         enumerant_attr_tokens.parse::<Token!(=)>()?;
 240         let enumerant: Ident = enumerant_attr_tokens.parse()?;
 241         input.parse::<Token!(fn)>()?;
 242         let fn_name: Ident = input.parse()?;
 243         let inputs_tokens;
 244         parenthesized!(inputs_tokens in input);
 245         let inputs = inputs_tokens.parse_terminated(InstructionInput::parse)?;
 246         input.parse::<Token!(->)>()?;
 247         let outputs_tokens;
 248         parenthesized!(outputs_tokens in input);
 249         let outputs = outputs_tokens.parse_terminated(InstructionOutput::parse)?;
 250         let body_tokens;
 251         braced!(body_tokens in input);
 252         let instruction_name: LitStr = body_tokens.parse()?;
 253         Ok(Self {
 254             enumerant,
 255             fn_name,
 256             inputs,
 257             outputs,
 258             instruction_name,
 259         })
 260     }
 261 }
 262
 263 impl Instruction {
 264     fn map_input_registers(&self) -> syn::Result<Vec<TokenStream>> {
 265         let mut retval = Vec::new();
 266         for input in &self.inputs {
 267             match input {
 268                 InstructionInput::Ra(_) => retval.push(quote! {InstructionInputRegister::Ra}),
 269                 InstructionInput::Rb(_) => retval.push(quote! {InstructionInputRegister::Rb}),
 270                 InstructionInput::Rc(_) => retval.push(quote! {InstructionInputRegister::Rc}),
 271                 InstructionInput::Carry(_) => retval.push(quote! {InstructionInputRegister::Carry}),
 272             }
 273         }
 274         Ok(retval)
 275     }
 276     fn to_native_fn_tokens(&self) -> syn::Result<TokenStream> {
 277         let Instruction {
 278             enumerant,
 279             fn_name,
 280             inputs,
 281             outputs,
 282             instruction_name,
 283         } = self;
 284         let mut asm = InlineAssembly::new(instruction_name.span());
 285         let mut before_asm = Vec::<TokenStream>::new();
 286         let mut after_asm = Vec::<TokenStream>::new();
 287         for output in &self.outputs {
 288             match output {
 289                 InstructionOutput::Rt(span) => {
 290                     unimplemented!("InstructionOutput::Rt");
 291                 }
 292                 InstructionOutput::Carry(span) => {
 293                     unimplemented!("InstructionOutput::Carry");
 294                 }
 295                 InstructionOutput::Overflow(span) => {
 296                     unimplemented!("InstructionOutput::Overflow");
 297                 }
 298                 InstructionOutput::CR0(span) => {
 299                     unimplemented!("InstructionOutput::CR0");
 300                 }
 301             }
 302         }
 303         for input in &self.inputs {
 304             match input {
 305                 InstructionInput::Ra(span) => {
 306                     unimplemented!("InstructionInput::Ra");
 307                 }
 308                 InstructionInput::Rb(span) => {
 309                     unimplemented!("InstructionInput::Rb");
 310                 }
 311                 InstructionInput::Rc(span) => {
 312                     unimplemented!("InstructionInput::Rc");
 313                 }
 314                 InstructionInput::Carry(span) => {
 315                     unimplemented!("InstructionInput::Carry");
 316                 }
 317             }
 318         }
 319         Ok(quote! {
 320             pub fn #fn_name(inputs: InstructionInput) -> InstructionResult {
 321                 #![allow(unused_variables, unused_assignments)]
 322                 let InstructionInput {
 323                     ra,
 324                     rb,
 325                     rc,
 326                     carry,
 327                 } = inputs;
 328                 let rt: u64;
 329                 let xer: u64;
 330                 let cr: u32;
 331                 #(#before_asm)*
 332                 unsafe {
 333                     #asm;
 334                 }
 335                 let mut retval = InstructionOutput::default();
 336                 #(#after_asm)*
 337                 retval
 338             }
 339         })
 340     }
 341 }
 342
 343 #[derive(Debug)]
 344 struct Instructions {
 345     instructions: Vec<Instruction>,
 346 }
 347
 348 impl Instructions {
 349     fn to_tokens(&self) -> syn::Result<TokenStream> {
 350         let mut fn_names = Vec::new();
 351         let mut instr_enumerants = Vec::new();
 352         let mut get_native_fn_match_cases = Vec::new();
 353         let mut get_model_fn_match_cases = Vec::new();
 354         let mut get_used_input_registers_match_cases = Vec::new();
 355         let mut name_match_cases = Vec::new();
 356         let mut enumerants = Vec::new();
 357         let mut native_fn_tokens = Vec::new();
 358         for instruction in &self.instructions {
 359             let Instruction {
 360                 enumerant,
 361                 fn_name,
 362                 inputs,
 363                 outputs,
 364                 instruction_name,
 365             } = instruction;
 366             fn_names.push(fn_name);
 367             enumerants.push(enumerant);
 368             instr_enumerants.push(quote! {
 369                 #[serde(rename = #instruction_name)]
 370                 #enumerant,
 371             });
 372             get_native_fn_match_cases.push(quote! {
 373                 Self::#enumerant => native_instrs::#fn_name,
 374             });
 375             get_model_fn_match_cases.push(quote! {
 376                 Self::#enumerant => instr_models::#fn_name,
 377             });
 378             let mapped_input_registers = instruction.map_input_registers()?;
 379             get_used_input_registers_match_cases.push(quote! {
 380                 Self::#enumerant => &[#(#mapped_input_registers),*],
 381             });
 382             name_match_cases.push(quote! {
 383                 Self::#enumerant => #instruction_name,
 384             });
 385             native_fn_tokens.push(instruction.to_native_fn_tokens()?);
 386         }
 387         Ok(quote! {
 388             #[cfg(feature = "python")]
 389             macro_rules! wrap_all_instr_fns {
 390                 ($m:ident) => {
 391                     wrap_instr_fns! {
 392                         #![pymodule($m)]
 393
 394                         #(fn #fn_names(inputs: InstructionInput) -> InstructionOutput;)*
 395                     }
 396                 };
 397             }
 398
 399             #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Serialize, Deserialize)]
 400             pub enum Instr {
 401                 #(#instr_enumerants)*
 402             }
 403
 404             impl Instr {
 405                 #[cfg(feature = "native_instrs")]
 406                 pub fn get_native_fn(self) -> fn(InstructionInput) -> InstructionOutput {
 407                     match self {
 408                         #(#get_native_fn_match_cases)*
 409                     }
 410                 }
 411                 pub fn get_model_fn(self) -> fn(InstructionInput) -> InstructionOutput {
 412                     match self {
 413                         #(#get_model_fn_match_cases)*
 414                     }
 415                 }
 416                 pub fn get_used_input_registers(self) -> &'static [InstructionInputRegister] {
 417                     match self {
 418                         #(#get_used_input_registers_match_cases)*
 419                     }
 420                 }
 421                 pub fn name(self) -> &'static str {
 422                     match self {
 423                         #(#name_match_cases)*
 424                     }
 425                 }
 426                 pub const VALUES: &'static [Self] = &[
 427                     #(Self::#enumerants,)*
 428                 ];
 429             }
 430
 431             #[cfg(feature = "native_instrs")]
 432             pub mod native_instrs {
 433                 use super::*;
 434
 435                 #(#native_fn_tokens)*
 436             }
 437         })
 438     }
 439 }
 440
 441 impl Parse for Instructions {
 442     fn parse(input: ParseStream) -> syn::Result<Self> {
 443         let mut instructions = Vec::new();
 444         while !input.is_empty() {
 445             instructions.push(input.parse()?);
 446         }
 447         Ok(Self { instructions })
 448     }
 449 }
 450
 451 #[proc_macro]
 452 pub fn instructions(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
 453     let input = parse_macro_input!(input as Instructions);
 454     match input.to_tokens() {
 455         Ok(retval) => retval,
 456         Err(err) => err.to_compile_error(),
 457     }
 458     .into()
 459 }