power-instruction-analyzer-proc-macro/src/inline_assembly.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::{format_ident, quote, quote_spanned, ToTokens};
   6 use std::{
   7     collections::HashMap,
   8     fmt::Write,
   9     hash::Hash,
  10     ops::{Deref, DerefMut},
  11     sync::atomic::{AtomicU64, Ordering},
  12 };
  13 use syn::{
  14     ext::IdentExt,
  15     parenthesized,
  16     parse::{Parse, ParseStream},
  17     punctuated::Punctuated,
  18     token::Paren,
  19     LitStr, Token,
  20 };
  21
  22 macro_rules! append_assembly {
  23     ($retval:ident;) => {};
  24     ($retval:ident; $lit:literal $($tt:tt)*) => {
  25         $crate::inline_assembly::ToAssembly::append_to($lit, &mut $retval);
  26         append_assembly!($retval; $($tt)*);
  27     };
  28     ($retval:ident; input($arg_id:ident = {$($arg_tt:tt)*}) $($tt:tt)*) => {
  29         {
  30             let (arg, arg_id) = $crate::inline_assembly::Assembly::make_input(quote! {$($arg_tt)*});
  31             $crate::inline_assembly::ToAssembly::append_to(&arg, &mut $retval);
  32             $arg_id = arg_id;
  33         }
  34         append_assembly!($retval; $($tt)*);
  35     };
  36     ($retval:ident; input{$($arg_tt:tt)*} $($tt:tt)*) => {
  37         {
  38             let (arg, _arg_id) = $crate::inline_assembly::Assembly::make_input(quote! {$($arg_tt)*});
  39             $crate::inline_assembly::ToAssembly::append_to(&arg, &mut $retval);
  40         }
  41         append_assembly!($retval; $($tt)*);
  42     };
  43     ($retval:ident; output($arg_id:ident = {$($arg_tt:tt)*}) $($tt:tt)*) => {
  44         {
  45             let (arg, arg_id) = $crate::inline_assembly::Assembly::make_output(quote! {$($arg_tt)*});
  46             $crate::inline_assembly::ToAssembly::append_to(&arg, &mut $retval);
  47             $arg_id = arg_id;
  48         }
  49         append_assembly!($retval; $($tt)*);
  50     };
  51     ($retval:ident; output{$($arg_tt:tt)*} $($tt:tt)*) => {
  52         {
  53             let (arg, _arg_id) = $crate::inline_assembly::Assembly::make_output(quote! {$($arg_tt)*});
  54             $crate::inline_assembly::ToAssembly::append_to(&arg, &mut $retval);
  55         }
  56         append_assembly!($retval; $($tt)*);
  57     };
  58     ($retval:ident; clobber{$($arg_tt:tt)*} $($tt:tt)*) => {
  59         $crate::inline_assembly::ToAssembly::append_to(
  60             &$crate::inline_assembly::Assembly::make_clobber(quote::quote! {$($arg_tt)*}),
  61             &mut $retval
  62         );
  63         append_assembly!($retval; $($tt)*);
  64     };
  65     ($retval:ident; ($arg_id:ident) $($tt:tt)*) => {
  66         $crate::inline_assembly::ToAssembly::append_to(&$arg_id, &mut $retval);
  67         append_assembly!($retval; $($tt)*);
  68     };
  69 }
  70
  71 macro_rules! assembly {
  72     () => {
  73         $crate::inline_assembly::Assembly::new()
  74     };
  75     ($($tt:tt)*) => {
  76         {
  77             let mut retval = $crate::inline_assembly::Assembly::new();
  78             append_assembly!(retval; $($tt)*);
  79             retval
  80         }
  81     };
  82 }
  83
  84 pub(crate) trait ToAssembly {
  85     /// appends `self` to `retval`
  86     fn append_to(&self, retval: &mut Assembly);
  87
  88     fn to_assembly(&self) -> Assembly {
  89         let mut retval = Assembly::default();
  90         self.append_to(&mut retval);
  91         retval
  92     }
  93
  94     fn into_assembly(self) -> Assembly
  95     where
  96         Self: Sized,
  97     {
  98         let mut retval = Assembly::default();
  99         self.append_to(&mut retval);
 100         retval
 101     }
 102 }
 103
 104 impl<T: ToAssembly + ?Sized> ToAssembly for &'_ T {
 105     fn append_to(&self, retval: &mut Assembly) {
 106         (**self).append_to(retval);
 107     }
 108
 109     fn to_assembly(&self) -> Assembly {
 110         (**self).to_assembly()
 111     }
 112 }
 113
 114 impl<T: ToAssembly + ?Sized> ToAssembly for &'_ mut T {
 115     fn append_to(&self, retval: &mut Assembly) {
 116         (**self).append_to(retval);
 117     }
 118
 119     fn to_assembly(&self) -> Assembly {
 120         (**self).to_assembly()
 121     }
 122 }
 123
 124 impl<T: ToAssembly> ToAssembly for Box<T> {
 125     fn append_to(&self, retval: &mut Assembly) {
 126         (**self).append_to(retval);
 127     }
 128
 129     fn to_assembly(&self) -> Assembly {
 130         (**self).to_assembly()
 131     }
 132
 133     fn into_assembly(self) -> Assembly {
 134         (*self).into_assembly()
 135     }
 136 }
 137
 138 impl ToAssembly for str {
 139     fn append_to(&self, retval: &mut Assembly) {
 140         if let Some(AssemblyTextFragment::Text(text)) = retval.text_fragments.last_mut() {
 141             *text += self;
 142         } else {
 143             retval
 144                 .text_fragments
 145                 .push(AssemblyTextFragment::Text(self.into()));
 146         }
 147     }
 148 }
 149
 150 impl ToAssembly for String {
 151     fn append_to(&self, retval: &mut Assembly) {
 152         str::append_to(&self, retval)
 153     }
 154 }
 155
 156 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 157 pub(crate) struct AssemblyMetavariableId(u64);
 158
 159 impl AssemblyMetavariableId {
 160     pub(crate) fn new() -> Self {
 161         // don't start at zero to help avoid confusing id with indexes
 162         static NEXT_ID: AtomicU64 = AtomicU64::new(10000);
 163         AssemblyMetavariableId(NEXT_ID.fetch_add(1, Ordering::Relaxed))
 164     }
 165 }
 166
 167 impl ToAssembly for AssemblyMetavariableId {
 168     fn append_to(&self, retval: &mut Assembly) {
 169         retval
 170             .text_fragments
 171             .push(AssemblyTextFragment::Metavariable(*self));
 172     }
 173 }
 174
 175 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 176 pub(crate) struct AssemblyArgId(u64);
 177
 178 impl AssemblyArgId {
 179     pub(crate) fn new() -> Self {
 180         // don't start at zero to help avoid confusing id with indexes
 181         static NEXT_ID: AtomicU64 = AtomicU64::new(1000);
 182         AssemblyArgId(NEXT_ID.fetch_add(1, Ordering::Relaxed))
 183     }
 184 }
 185
 186 impl ToAssembly for AssemblyArgId {
 187     fn append_to(&self, retval: &mut Assembly) {
 188         retval
 189             .text_fragments
 190             .push(AssemblyTextFragment::ArgIndex(*self));
 191     }
 192 }
 193
 194 macro_rules! impl_assembly_arg {
 195     (
 196         struct $name:ident {
 197             tokens: TokenStream,
 198             $(
 199                 $id:ident: AssemblyArgId,
 200             )?
 201         }
 202     ) => {
 203         #[derive(Debug, Clone)]
 204         struct $name {
 205             tokens: TokenStream,
 206             $($id: AssemblyArgId,)?
 207         }
 208
 209         impl ToTokens for $name {
 210             fn to_token_stream(&self) -> TokenStream {
 211                 self.tokens.clone()
 212             }
 213
 214             fn into_token_stream(self) -> TokenStream {
 215                 self.tokens
 216             }
 217
 218             fn to_tokens(&self, tokens: &mut TokenStream) {
 219                 self.tokens.to_tokens(tokens)
 220             }
 221         }
 222
 223         impl From<TokenStream> for $name {
 224             fn from(tokens: TokenStream) -> Self {
 225                 Self {
 226                     tokens,
 227                     $($id: AssemblyArgId::new(),)?
 228                 }
 229             }
 230         }
 231     };
 232 }
 233
 234 impl_assembly_arg! {
 235     struct AssemblyInputArg {
 236         tokens: TokenStream,
 237         id: AssemblyArgId,
 238     }
 239 }
 240
 241 impl_assembly_arg! {
 242     struct AssemblyOutputArg {
 243         tokens: TokenStream,
 244         id: AssemblyArgId,
 245     }
 246 }
 247
 248 impl_assembly_arg! {
 249     struct AssemblyClobber {
 250         tokens: TokenStream,
 251     }
 252 }
 253
 254 #[derive(Debug, Clone)]
 255 pub(crate) enum AssemblyTextFragment {
 256     Text(String),
 257     ArgIndex(AssemblyArgId),
 258     Metavariable(AssemblyMetavariableId),
 259 }
 260
 261 #[derive(Debug, Default, Clone)]
 262 pub(crate) struct Assembly {
 263     text_fragments: Vec<AssemblyTextFragment>,
 264     inputs: Vec<AssemblyInputArg>,
 265     outputs: Vec<AssemblyOutputArg>,
 266     clobbers: Vec<AssemblyClobber>,
 267 }
 268
 269 impl From<String> for Assembly {
 270     fn from(text: String) -> Self {
 271         Self {
 272             text_fragments: vec![AssemblyTextFragment::Text(text)],
 273             ..Self::default()
 274         }
 275     }
 276 }
 277
 278 impl From<&'_ str> for Assembly {
 279     fn from(text: &str) -> Self {
 280         String::from(text).into()
 281     }
 282 }
 283
 284 impl From<AssemblyArgId> for Assembly {
 285     fn from(arg_id: AssemblyArgId) -> Self {
 286         Self {
 287             text_fragments: vec![AssemblyTextFragment::ArgIndex(arg_id)],
 288             ..Self::default()
 289         }
 290     }
 291 }
 292
 293 impl From<&'_ AssemblyArgId> for Assembly {
 294     fn from(arg_id: &AssemblyArgId) -> Self {
 295         Self::from(*arg_id)
 296     }
 297 }
 298
 299 impl From<AssemblyMetavariableId> for Assembly {
 300     fn from(arg_id: AssemblyMetavariableId) -> Self {
 301         Self {
 302             text_fragments: vec![AssemblyTextFragment::Metavariable(arg_id)],
 303             ..Self::default()
 304         }
 305     }
 306 }
 307
 308 impl From<&'_ AssemblyMetavariableId> for Assembly {
 309     fn from(arg_id: &AssemblyMetavariableId) -> Self {
 310         Self::from(*arg_id)
 311     }
 312 }
 313
 314 impl Assembly {
 315     pub(crate) fn new() -> Self {
 316         Self::default()
 317     }
 318     pub(crate) fn make_input(tokens: impl ToTokens) -> (Self, AssemblyArgId) {
 319         let input: AssemblyInputArg = tokens.into_token_stream().into();
 320         let id = input.id;
 321         (
 322             Self {
 323                 text_fragments: vec![AssemblyTextFragment::ArgIndex(id)],
 324                 inputs: vec![input],
 325                 ..Self::default()
 326             },
 327             id,
 328         )
 329     }
 330     pub(crate) fn make_output(tokens: impl ToTokens) -> (Self, AssemblyArgId) {
 331         let output: AssemblyOutputArg = tokens.into_token_stream().into();
 332         let id = output.id;
 333         (
 334             Self {
 335                 text_fragments: vec![AssemblyTextFragment::ArgIndex(id)],
 336                 outputs: vec![output],
 337                 ..Self::default()
 338             },
 339             id,
 340         )
 341     }
 342     pub(crate) fn make_clobber(tokens: impl ToTokens) -> Self {
 343         Self {
 344             clobbers: vec![tokens.into_token_stream().into()],
 345             ..Self::default()
 346         }
 347     }
 348     pub(crate) fn replace_metavariables<R>(
 349         &self,
 350         mut f: impl FnMut(AssemblyMetavariableId) -> Result<Assembly, R>,
 351     ) -> Result<Assembly, R> {
 352         let mut retval = self.args_without_text();
 353         for text_fragment in &self.text_fragments {
 354             match text_fragment {
 355                 AssemblyTextFragment::Text(text) => text.append_to(&mut retval),
 356                 AssemblyTextFragment::ArgIndex(id) => id.append_to(&mut retval),
 357                 AssemblyTextFragment::Metavariable(id) => f(*id)?.append_to(&mut retval),
 358             }
 359         }
 360         Ok(retval)
 361     }
 362     pub(crate) fn args_without_text(&self) -> Assembly {
 363         Assembly {
 364             text_fragments: Vec::new(),
 365             inputs: self.inputs.clone(),
 366             outputs: self.outputs.clone(),
 367             clobbers: self.clobbers.clone(),
 368         }
 369     }
 370     pub(crate) fn text_without_args(&self) -> Assembly {
 371         Assembly {
 372             text_fragments: self.text_fragments.clone(),
 373             inputs: Vec::new(),
 374             outputs: Vec::new(),
 375             clobbers: Vec::new(),
 376         }
 377     }
 378     pub(crate) fn to_text(&self) -> String {
 379         let mut id_index_map = HashMap::new();
 380         for (index, id) in self
 381             .outputs
 382             .iter()
 383             .map(|v| v.id)
 384             .chain(self.inputs.iter().map(|v| v.id))
 385             .enumerate()
 386         {
 387             if let Some(old_index) = id_index_map.insert(id, index) {
 388                 panic!(
 389                     "duplicate id in inline assembly arguments: #{} and #{}\n{:#?}",
 390                     old_index, index, self
 391                 );
 392             }
 393         }
 394         let mut retval = String::new();
 395         for text_fragment in &self.text_fragments {
 396             match text_fragment {
 397                 AssemblyTextFragment::Text(text) => retval += text,
 398                 AssemblyTextFragment::Metavariable(id) => {
 399                     panic!(
 400                         "metavariables are not allowed when converting \
 401                             assembly to text: metavariable id={:?}\n{:#?}",
 402                         id, self
 403                     );
 404                 }
 405                 AssemblyTextFragment::ArgIndex(id) => {
 406                     if let Some(index) = id_index_map.get(id) {
 407                         write!(retval, "arg{}", index).unwrap();
 408                     } else {
 409                         panic!(
 410                             "unknown id in inline assembly arguments: id={:?}\n{:#?}",
 411                             id, self
 412                         );
 413                     }
 414                 }
 415             }
 416         }
 417         retval
 418     }
 419 }
 420
 421 impl ToAssembly for Assembly {
 422     fn append_to(&self, retval: &mut Assembly) {
 423         retval.text_fragments.reserve(self.text_fragments.len());
 424         for text_fragment in &self.text_fragments {
 425             match *text_fragment {
 426                 AssemblyTextFragment::Text(ref text) => text.append_to(retval),
 427                 AssemblyTextFragment::Metavariable(id) => id.append_to(retval),
 428                 AssemblyTextFragment::ArgIndex(id) => id.append_to(retval),
 429             }
 430         }
 431         retval.inputs.extend_from_slice(&self.inputs);
 432         retval.outputs.extend_from_slice(&self.outputs);
 433         retval.clobbers.extend_from_slice(&self.clobbers);
 434     }
 435
 436     fn to_assembly(&self) -> Assembly {
 437         self.clone()
 438     }
 439
 440     fn into_assembly(self) -> Assembly {
 441         self
 442     }
 443 }
 444
 445 #[derive(Debug, Clone)]
 446 pub(crate) struct AssemblyWithTextSpan {
 447     pub(crate) asm: Assembly,
 448     pub(crate) text_span: Span,
 449 }
 450
 451 impl Deref for AssemblyWithTextSpan {
 452     type Target = Assembly;
 453
 454     fn deref(&self) -> &Self::Target {
 455         &self.asm
 456     }
 457 }
 458
 459 impl DerefMut for AssemblyWithTextSpan {
 460     fn deref_mut(&mut self) -> &mut Self::Target {
 461         &mut self.asm
 462     }
 463 }
 464
 465 impl ToTokens for AssemblyWithTextSpan {
 466     fn to_tokens(&self, tokens: &mut TokenStream) {
 467         let Self {
 468             asm:
 469                 Assembly {
 470                     text_fragments: _,
 471                     inputs,
 472                     outputs,
 473                     clobbers,
 474                 },
 475             text_span,
 476         } = self;
 477         let mut args: Punctuated<TokenStream, Token![,]> = self
 478             .to_text()
 479             .lines()
 480             .map(|line| {
 481                 quote_spanned! {*text_span=>
 482                     #line
 483                 }
 484             })
 485             .collect();
 486         let mut named_args = Vec::new();
 487         let mut unnamed_args = Vec::new();
 488         for (index, tokens) in outputs
 489             .iter()
 490             .map(|v| &v.tokens)
 491             .chain(inputs.iter().map(|v| &v.tokens))
 492             .enumerate()
 493         {
 494             match syn::parse2::<AsmArg>(tokens.clone())
 495                 .unwrap_or_else(|e| panic!("failed to parse AsmArg: {}\nTokens:\n{}", e, tokens))
 496                 .reg
 497             {
 498                 AsmArgReg::RegClass(_) => {
 499                     let id = format_ident!("arg{}", index);
 500                     named_args.push(quote! { #id = #tokens });
 501                 }
 502                 AsmArgReg::RegLit(_) => unnamed_args.push(tokens.clone()),
 503             }
 504         }
 505         args.extend(named_args);
 506         args.extend(unnamed_args);
 507         args.extend(clobbers.iter().map(ToTokens::to_token_stream));
 508         let value = quote! {
 509             asm!(#args)
 510         };
 511         value.to_tokens(tokens);
 512     }
 513 }
 514
 515 enum AsmArgReg {
 516     RegClass(Ident),
 517     RegLit(LitStr),
 518 }
 519
 520 impl Parse for AsmArgReg {
 521     fn parse(input: ParseStream) -> syn::Result<Self> {
 522         if input.peek(Ident::peek_any) {
 523             Ok(Self::RegClass(input.call(Ident::parse_any)?))
 524         } else {
 525             Ok(Self::RegLit(input.parse()?))
 526         }
 527     }
 528 }
 529
 530 #[allow(dead_code)]
 531 struct AsmArg {
 532     io_kind: Ident,
 533     paren: Paren,
 534     reg: AsmArgReg,
 535     body: TokenStream,
 536 }
 537
 538 impl Parse for AsmArg {
 539     fn parse(input: ParseStream) -> syn::Result<Self> {
 540         let input2;
 541         Ok(Self {
 542             io_kind: input.call(Ident::parse_any)?,
 543             paren: parenthesized!(input2 in input),
 544             reg: input2.parse()?,
 545             body: input.parse()?,
 546         })
 547     }
 548 }