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, ToTokens};
   6 use std::{
   7     borrow::Cow,
   8     collections::HashMap,
   9     fmt::Write,
  10     hash::{Hash, Hasher},
  11     marker::PhantomPinned,
  12     ops::{Add, AddAssign, Deref, DerefMut},
  13     pin::Pin,
  14     rc::Rc,
  15     sync::atomic::{AtomicU64, Ordering},
  16 };
  17 use syn::LitStr;
  18
  19 pub(crate) trait ToAssembly {
  20     fn to_assembly(&self) -> Assembly;
  21 }
  22
  23 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
  24 pub(crate) struct AssemblyArgId(u64);
  25
  26 impl AssemblyArgId {
  27     pub(crate) fn new() -> Self {
  28         // don't start at zero to help avoid confusing id with indexes
  29         static NEXT_ID: AtomicU64 = AtomicU64::new(1000);
  30         AssemblyArgId(NEXT_ID.fetch_add(1, Ordering::Relaxed))
  31     }
  32 }
  33
  34 macro_rules! impl_assembly_arg {
  35     (
  36         $vis:vis struct $name:ident {
  37             tokens: TokenStream,
  38             $(
  39                 $id:ident: AssemblyArgId,
  40             )?
  41         }
  42     ) => {
  43         #[derive(Debug, Clone)]
  44         $vis struct $name {
  45             tokens: TokenStream,
  46             $($id: AssemblyArgId,)?
  47         }
  48
  49         impl $name {
  50             $vis fn new(tokens: impl ToTokens) -> Self {
  51                 tokens.into_token_stream().into()
  52             }
  53         }
  54
  55         impl ToTokens for $name {
  56             fn to_token_stream(&self) -> TokenStream {
  57                 self.tokens.clone()
  58             }
  59
  60             fn into_token_stream(self) -> TokenStream {
  61                 self.tokens
  62             }
  63
  64             fn to_tokens(&self, tokens: &mut TokenStream) {
  65                 self.tokens.to_tokens(tokens)
  66             }
  67         }
  68
  69         impl From<TokenStream> for $name {
  70             fn from(tokens: TokenStream) -> Self {
  71                 Self {
  72                     tokens,
  73                     $($id: AssemblyArgId::new(),)?
  74                 }
  75             }
  76         }
  77     };
  78 }
  79
  80 impl_assembly_arg! {
  81     pub(crate) struct AssemblyInputArg {
  82         tokens: TokenStream,
  83         id: AssemblyArgId,
  84     }
  85 }
  86
  87 impl_assembly_arg! {
  88     pub(crate) struct AssemblyOutputArg {
  89         tokens: TokenStream,
  90         id: AssemblyArgId,
  91     }
  92 }
  93
  94 impl_assembly_arg! {
  95     pub(crate) struct AssemblyClobber {
  96         tokens: TokenStream,
  97     }
  98 }
  99
 100 #[derive(Debug, Clone)]
 101 pub(crate) enum AssemblyTextFragment {
 102     Text(String),
 103     ArgIndex(AssemblyArgId),
 104 }
 105
 106 #[derive(Debug, Default, Clone)]
 107 pub(crate) struct Assembly {
 108     text_fragments: Vec<AssemblyTextFragment>,
 109     inputs: Vec<AssemblyInputArg>,
 110     outputs: Vec<AssemblyOutputArg>,
 111     clobbers: Vec<AssemblyClobber>,
 112 }
 113
 114 impl From<String> for Assembly {
 115     fn from(text: String) -> Self {
 116         Self {
 117             text_fragments: vec![AssemblyTextFragment::Text(text)],
 118             ..Self::default()
 119         }
 120     }
 121 }
 122
 123 impl From<&'_ str> for Assembly {
 124     fn from(text: &str) -> Self {
 125         String::from(text).into()
 126     }
 127 }
 128
 129 impl Assembly {
 130     pub(crate) fn new() -> Self {
 131         Self::default()
 132     }
 133     pub(crate) fn to_text(&self) -> String {
 134         let mut id_index_map = HashMap::new();
 135         for (index, id) in self
 136             .outputs
 137             .iter()
 138             .map(|v| v.id)
 139             .chain(self.inputs.iter().map(|v| v.id))
 140             .enumerate()
 141         {
 142             if let Some(old_index) = id_index_map.insert(id, index) {
 143                 panic!(
 144                     "duplicate id in inline assembly arguments: #{} and #{}\n{:#?}",
 145                     old_index, index, self
 146                 );
 147             }
 148         }
 149         let mut retval = String::new();
 150         for text_fragment in &self.text_fragments {
 151             match text_fragment {
 152                 AssemblyTextFragment::Text(text) => retval += text,
 153                 AssemblyTextFragment::ArgIndex(id) => {
 154                     if let Some(index) = id_index_map.get(id) {
 155                         write!(retval, "{}", index).unwrap();
 156                     } else {
 157                         panic!(
 158                             "unknown id in inline assembly arguments: id={:?}\n{:#?}",
 159                             id, self
 160                         );
 161                     }
 162                 }
 163             }
 164         }
 165         retval
 166     }
 167 }
 168
 169 impl AddAssign<&'_ Assembly> for Assembly {
 170     fn add_assign(&mut self, rhs: &Assembly) {
 171         let Self {
 172             text_fragments,
 173             inputs,
 174             outputs,
 175             clobbers,
 176         } = self;
 177         text_fragments.reserve(rhs.text_fragments.len());
 178         for text_fragment in &rhs.text_fragments {
 179             match *text_fragment {
 180                 AssemblyTextFragment::Text(ref rhs_text) => {
 181                     if let Some(AssemblyTextFragment::Text(text)) = text_fragments.last_mut() {
 182                         *text += rhs_text;
 183                     } else {
 184                         text_fragments.push(AssemblyTextFragment::Text(rhs_text.clone()));
 185                     }
 186                 }
 187                 AssemblyTextFragment::ArgIndex(id) => {
 188                     self.text_fragments.push(AssemblyTextFragment::ArgIndex(id));
 189                 }
 190             }
 191         }
 192         inputs.extend_from_slice(&rhs.inputs);
 193         outputs.extend_from_slice(&rhs.outputs);
 194         clobbers.extend_from_slice(&rhs.clobbers);
 195     }
 196 }
 197
 198 impl AddAssign<Assembly> for Assembly {
 199     fn add_assign(&mut self, rhs: Assembly) {
 200         *self += &rhs;
 201     }
 202 }
 203
 204 impl Add for Assembly {
 205     type Output = Assembly;
 206
 207     fn add(mut self, rhs: Self) -> Self::Output {
 208         self += rhs;
 209         self
 210     }
 211 }
 212
 213 impl Add<&'_ Assembly> for Assembly {
 214     type Output = Assembly;
 215
 216     fn add(mut self, rhs: &Assembly) -> Self::Output {
 217         self += rhs;
 218         self
 219     }
 220 }
 221
 222 impl Add<Assembly> for &'_ Assembly {
 223     type Output = Assembly;
 224
 225     fn add(self, rhs: Assembly) -> Self::Output {
 226         Assembly::clone(self) + rhs
 227     }
 228 }
 229
 230 impl Add<&'_ Assembly> for &'_ Assembly {
 231     type Output = Assembly;
 232
 233     fn add(self, rhs: &Assembly) -> Self::Output {
 234         Assembly::clone(self) + rhs
 235     }
 236 }
 237
 238 #[derive(Debug, Clone)]
 239 pub(crate) struct AssemblyWithTextSpan {
 240     pub(crate) asm: Assembly,
 241     pub(crate) text_span: Span,
 242 }
 243
 244 impl Deref for AssemblyWithTextSpan {
 245     type Target = Assembly;
 246
 247     fn deref(&self) -> &Self::Target {
 248         &self.asm
 249     }
 250 }
 251
 252 impl DerefMut for AssemblyWithTextSpan {
 253     fn deref_mut(&mut self) -> &mut Self::Target {
 254         &mut self.asm
 255     }
 256 }
 257
 258 impl ToTokens for AssemblyWithTextSpan {
 259     fn to_tokens(&self, tokens: &mut TokenStream) {
 260         let Self {
 261             asm:
 262                 Assembly {
 263                     text_fragments: _,
 264                     inputs,
 265                     outputs,
 266                     clobbers,
 267                 },
 268             text_span,
 269         } = self;
 270         let text = LitStr::new(&self.to_text(), text_span.clone());
 271         let value = quote! {
 272             llvm_asm!(#text : #(#outputs),* : #(#inputs),* : #(#clobbers),*)
 273         };
 274         value.to_tokens(tokens);
 275     }
 276 }