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