register_allocator/src/index.rs

   1 use serde::{de::DeserializeOwned, Deserialize, Serialize};
   2 use std::{
   3     fmt,
   4     hash::Hash,
   5     iter::FusedIterator,
   6     ops::{Index, IndexMut, Range},
   7 };
   8
   9 use crate::{
  10     error::{
  11         BlockIdxOutOfRange, BlockParamIdxOutOfRange, InstIdxOutOfRange, OperandIdxOutOfRange,
  12         SSAValIdxOutOfRange,
  13     },
  14     function::{Block, Inst, Operand, SSAVal},
  15     live_range::LiveRange,
  16 };
  17
  18 pub trait DisplayOptionIdx {
  19     type Type: fmt::Display;
  20     fn display_option_idx(self) -> Self::Type;
  21 }
  22
  23 pub trait IndexTy: Copy + fmt::Debug + Ord + Hash + Serialize + DeserializeOwned {
  24     fn new(value: usize) -> Self;
  25     fn get(self) -> usize;
  26     fn next(self) -> Self {
  27         Self::new(self.get() + 1)
  28     }
  29     fn prev(self) -> Self {
  30         Self::new(self.get() - 1)
  31     }
  32 }
  33
  34 #[derive(Debug, Clone)]
  35 pub struct RangeIter<T: IndexTy>(pub Range<T>);
  36
  37 impl<T: IndexTy> RangeIter<T> {
  38     pub fn as_usize_range(&self) -> Range<usize> {
  39         self.0.start.get()..self.0.end.get()
  40     }
  41     pub fn from_usize_range(range: Range<usize>) -> Self {
  42         Self(T::new(range.start)..T::new(range.end))
  43     }
  44     pub fn with_self_as_usize_range<R, F: FnOnce(&mut Range<usize>) -> R>(&mut self, f: F) -> R {
  45         let mut range = self.as_usize_range();
  46         let retval = f(&mut range);
  47         *self = Self::from_usize_range(range);
  48         retval
  49     }
  50 }
  51
  52 impl<T: IndexTy> Iterator for RangeIter<T> {
  53     type Item = T;
  54
  55     fn next(&mut self) -> Option<Self::Item> {
  56         self.with_self_as_usize_range(|r| r.next().map(T::new))
  57     }
  58
  59     fn size_hint(&self) -> (usize, Option<usize>) {
  60         self.as_usize_range().size_hint()
  61     }
  62
  63     fn nth(&mut self, n: usize) -> Option<Self::Item> {
  64         self.with_self_as_usize_range(|r| r.nth(n).map(T::new))
  65     }
  66
  67     fn fold<B, F>(self, init: B, mut f: F) -> B
  68     where
  69         F: FnMut(B, Self::Item) -> B,
  70     {
  71         self.as_usize_range()
  72             .fold(init, move |a, v| f(a, T::new(v)))
  73     }
  74 }
  75
  76 impl<T: IndexTy> FusedIterator for RangeIter<T> {}
  77
  78 impl<T: IndexTy> ExactSizeIterator for RangeIter<T> {
  79     fn len(&self) -> usize {
  80         self.as_usize_range().len()
  81     }
  82 }
  83
  84 impl<T: IndexTy> DoubleEndedIterator for RangeIter<T> {
  85     fn next_back(&mut self) -> Option<Self::Item> {
  86         self.with_self_as_usize_range(|r| r.next_back().map(T::new))
  87     }
  88
  89     fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
  90         self.with_self_as_usize_range(|r| r.nth_back(n).map(T::new))
  91     }
  92
  93     fn rfold<B, F>(self, init: B, mut f: F) -> B
  94     where
  95         F: FnMut(B, Self::Item) -> B,
  96     {
  97         self.as_usize_range()
  98             .rfold(init, move |a, v| f(a, T::new(v)))
  99     }
 100 }
 101
 102 macro_rules! define_index {
 103     ($name:ident) => {
 104         #[derive(
 105             Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
 106         )]
 107         #[serde(transparent)]
 108         pub struct $name {
 109             value: usize,
 110         }
 111
 112         impl IndexTy for $name {
 113             fn new(value: usize) -> Self {
 114                 Self { value }
 115             }
 116             fn get(self) -> usize {
 117                 self.value
 118             }
 119         }
 120
 121         impl $name {
 122             pub const fn new(value: usize) -> Self {
 123                 Self { value }
 124             }
 125             pub const fn get(self) -> usize {
 126                 self.value
 127             }
 128             pub const fn next(self) -> Self {
 129                 Self {
 130                     value: self.value + 1,
 131                 }
 132             }
 133             pub const fn prev(self) -> Self {
 134                 Self {
 135                     value: self.value - 1,
 136                 }
 137             }
 138         }
 139
 140         const _: () = {
 141             #[derive(Copy, Clone)]
 142             pub struct DisplayOptionIdxImpl(Option<$name>);
 143
 144             impl fmt::Debug for DisplayOptionIdxImpl {
 145                 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 146                     fmt::Display::fmt(self, f)
 147                 }
 148             }
 149
 150             impl fmt::Display for DisplayOptionIdxImpl {
 151                 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 152                     match &self.0 {
 153                         Some(v) => fmt::Display::fmt(v, f),
 154                         None => write!(f, "none"),
 155                     }
 156                 }
 157             }
 158
 159             impl DisplayOptionIdx for Option<$name> {
 160                 type Type = DisplayOptionIdxImpl;
 161
 162                 fn display_option_idx(self) -> Self::Type {
 163                     DisplayOptionIdxImpl(self)
 164                 }
 165             }
 166         };
 167     };
 168 }
 169
 170 define_index!(SSAValIdx);
 171 define_index!(InstIdx);
 172 define_index!(BlockIdx);
 173 define_index!(BlockParamIdx);
 174 define_index!(OperandIdx);
 175 define_index!(LiveRangeIdx);
 176
 177 pub trait Entries<'a, I>: Index<I>
 178 where
 179     I: IndexTy,
 180     Self::Output: 'a,
 181 {
 182     type Iter: Iterator<Item = (I, &'a Self::Output)>
 183         + DoubleEndedIterator
 184         + ExactSizeIterator
 185         + FusedIterator;
 186     fn entries(&'a self) -> Self::Iter;
 187     fn keys(&'a self) -> RangeIter<I>;
 188 }
 189
 190 pub trait EntriesMut<'a, I>: Entries<'a, I> + IndexMut<I>
 191 where
 192     I: IndexTy,
 193     Self::Output: 'a,
 194 {
 195     type IterMut: Iterator<Item = (I, &'a mut Self::Output)>
 196         + DoubleEndedIterator
 197         + ExactSizeIterator
 198         + FusedIterator;
 199     fn entries_mut(&'a mut self) -> Self::IterMut;
 200 }
 201
 202 pub trait TryIndex<I>: for<'a> Entries<'a, I>
 203 where
 204     I: IndexTy,
 205 {
 206     type Error;
 207     fn try_index(&self, idx: I) -> Result<&Self::Output, Self::Error>;
 208 }
 209
 210 pub trait TryIndexMut<I>: TryIndex<I> + for<'a> EntriesMut<'a, I>
 211 where
 212     I: IndexTy,
 213 {
 214     fn try_index_mut(&mut self, idx: I) -> Result<&mut Self::Output, Self::Error>;
 215 }
 216
 217 macro_rules! impl_index {
 218     (
 219         #[$(error = $Error:ident, )?iter = $Iter:ident, iter_mut = $IterMut:ident]
 220         impl Index<$I:ty> for Vec<$T:ty> {}
 221     ) => {
 222         #[derive(Clone, Debug)]
 223         pub struct $Iter<'a> {
 224             iter: std::iter::Enumerate<std::slice::Iter<'a, $T>>,
 225         }
 226
 227         impl<'a> Iterator for $Iter<'a> {
 228             type Item = ($I, &'a $T);
 229
 230             fn next(&mut self) -> Option<Self::Item> {
 231                 self.iter.next().map(|(i, v)| (<$I>::new(i), v))
 232             }
 233
 234             fn size_hint(&self) -> (usize, Option<usize>) {
 235                 self.iter.size_hint()
 236             }
 237
 238             fn fold<B, F>(self, init: B, mut f: F) -> B
 239             where
 240                 F: FnMut(B, Self::Item) -> B,
 241             {
 242                 self.iter
 243                     .fold(init, move |a, (i, v)| f(a, (<$I>::new(i), v)))
 244             }
 245         }
 246
 247         impl DoubleEndedIterator for $Iter<'_> {
 248             fn next_back(&mut self) -> Option<Self::Item> {
 249                 self.iter.next_back().map(|(i, v)| (<$I>::new(i), v))
 250             }
 251
 252             fn rfold<B, F>(self, init: B, mut f: F) -> B
 253             where
 254                 F: FnMut(B, Self::Item) -> B,
 255             {
 256                 self.iter
 257                     .rfold(init, move |a, (i, v)| f(a, (<$I>::new(i), v)))
 258             }
 259         }
 260
 261         impl ExactSizeIterator for $Iter<'_> {
 262             fn len(&self) -> usize {
 263                 self.iter.len()
 264             }
 265         }
 266
 267         impl FusedIterator for $Iter<'_> {}
 268
 269         #[derive(Debug)]
 270         pub struct $IterMut<'a> {
 271             iter: std::iter::Enumerate<std::slice::IterMut<'a, $T>>,
 272         }
 273
 274         impl<'a> Iterator for $IterMut<'a> {
 275             type Item = ($I, &'a mut $T);
 276
 277             fn next(&mut self) -> Option<Self::Item> {
 278                 self.iter.next().map(|(i, v)| (<$I>::new(i), v))
 279             }
 280
 281             fn size_hint(&self) -> (usize, Option<usize>) {
 282                 self.iter.size_hint()
 283             }
 284
 285             fn fold<B, F>(self, init: B, mut f: F) -> B
 286             where
 287                 F: FnMut(B, Self::Item) -> B,
 288             {
 289                 self.iter
 290                     .fold(init, move |a, (i, v)| f(a, (<$I>::new(i), v)))
 291             }
 292         }
 293
 294         impl DoubleEndedIterator for $IterMut<'_> {
 295             fn next_back(&mut self) -> Option<Self::Item> {
 296                 self.iter.next_back().map(|(i, v)| (<$I>::new(i), v))
 297             }
 298
 299             fn rfold<B, F>(self, init: B, mut f: F) -> B
 300             where
 301                 F: FnMut(B, Self::Item) -> B,
 302             {
 303                 self.iter
 304                     .rfold(init, move |a, (i, v)| f(a, (<$I>::new(i), v)))
 305             }
 306         }
 307
 308         impl ExactSizeIterator for $IterMut<'_> {
 309             fn len(&self) -> usize {
 310                 self.iter.len()
 311             }
 312         }
 313
 314         impl FusedIterator for $IterMut<'_> {}
 315
 316         impl Index<$I> for Vec<$T> {
 317             type Output = $T;
 318
 319             fn index(&self, index: $I) -> &Self::Output {
 320                 &self[index.get()]
 321             }
 322         }
 323
 324         impl IndexMut<$I> for Vec<$T> {
 325             fn index_mut(&mut self, index: $I) -> &mut Self::Output {
 326                 &mut self[index.get()]
 327             }
 328         }
 329
 330         impl<'a> Entries<'a, $I> for Vec<$T> {
 331             type Iter = $Iter<'a>;
 332             fn entries(&'a self) -> Self::Iter {
 333                 $Iter {
 334                     iter: (**self).iter().enumerate(),
 335                 }
 336             }
 337             fn keys(&'a self) -> RangeIter<$I> {
 338                 RangeIter::from_usize_range(0..self.len())
 339             }
 340         }
 341
 342         impl<'a> EntriesMut<'a, $I> for Vec<$T> {
 343             type IterMut = $IterMut<'a>;
 344             fn entries_mut(&'a mut self) -> Self::IterMut {
 345                 $IterMut {
 346                     iter: (**self).iter_mut().enumerate(),
 347                 }
 348             }
 349         }
 350
 351         $(
 352             impl TryIndex<$I> for Vec<$T> {
 353                 type Error = $Error;
 354
 355                 fn try_index(&self, idx: $I) -> Result<&Self::Output, Self::Error> {
 356                     self.get(idx.get()).ok_or($Error { idx })
 357                 }
 358             }
 359
 360             impl TryIndexMut<$I> for Vec<$T> {
 361                 fn try_index_mut(&mut self, idx: $I) -> Result<&mut Self::Output, Self::Error> {
 362                     self.get_mut(idx.get()).ok_or($Error { idx })
 363                 }
 364             }
 365         )?
 366
 367         impl Index<$I> for [$T] {
 368             type Output = $T;
 369
 370             fn index(&self, index: $I) -> &Self::Output {
 371                 &self[index.get()]
 372             }
 373         }
 374
 375         impl IndexMut<$I> for [$T] {
 376             fn index_mut(&mut self, index: $I) -> &mut Self::Output {
 377                 &mut self[index.get()]
 378             }
 379         }
 380
 381         impl<'a> Entries<'a, $I> for [$T] {
 382             type Iter = $Iter<'a>;
 383             fn entries(&'a self) -> Self::Iter {
 384                 $Iter {
 385                     iter: self.iter().enumerate(),
 386                 }
 387             }
 388             fn keys(&'a self) -> RangeIter<$I> {
 389                 RangeIter::from_usize_range(0..self.len())
 390             }
 391         }
 392
 393         impl<'a> EntriesMut<'a, $I> for [$T] {
 394             type IterMut = $IterMut<'a>;
 395             fn entries_mut(&'a mut self) -> Self::IterMut {
 396                 $IterMut {
 397                     iter: self.iter_mut().enumerate(),
 398                 }
 399             }
 400         }
 401
 402         $(
 403             impl TryIndex<$I> for [$T] {
 404                 type Error = $Error;
 405
 406                 fn try_index(&self, idx: $I) -> Result<&Self::Output, Self::Error> {
 407                     self.get(idx.get()).ok_or($Error { idx })
 408                 }
 409             }
 410
 411             impl TryIndexMut<$I> for [$T] {
 412                 fn try_index_mut(&mut self, idx: $I) -> Result<&mut Self::Output, Self::Error> {
 413                     self.get_mut(idx.get()).ok_or($Error { idx })
 414                 }
 415             }
 416         )?
 417     };
 418 }
 419
 420 impl_index! {
 421     #[error = SSAValIdxOutOfRange, iter = SSAValEntriesIter, iter_mut = SSAValEntriesIterMut]
 422     impl Index<SSAValIdx> for Vec<SSAVal> {}
 423 }
 424
 425 impl_index! {
 426     #[error = BlockParamIdxOutOfRange, iter = BlockParamEntriesIter, iter_mut = BlockParamEntriesIterMut]
 427     impl Index<BlockParamIdx> for Vec<SSAValIdx> {}
 428 }
 429
 430 impl_index! {
 431     #[error = OperandIdxOutOfRange, iter = OperandEntriesIter, iter_mut = OperandEntriesIterMut]
 432     impl Index<OperandIdx> for Vec<Operand> {}
 433 }
 434
 435 impl_index! {
 436     #[error = InstIdxOutOfRange, iter = InstEntriesIter, iter_mut = InstEntriesIterMut]
 437     impl Index<InstIdx> for Vec<Inst> {}
 438 }
 439
 440 impl_index! {
 441     #[error = BlockIdxOutOfRange, iter = BlockEntriesIter, iter_mut = BlockEntriesIterMut]
 442     impl Index<BlockIdx> for Vec<Block> {}
 443 }
 444
 445 impl_index! {
 446     #[iter = LiveRangeEntriesIter, iter_mut = LiveRangeEntriesIterMut]
 447     impl Index<LiveRangeIdx> for Vec<LiveRange> {}
 448 }
 449
 450 impl fmt::Display for SSAValIdx {
 451     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 452         write!(f, "v{}", self.get())
 453     }
 454 }
 455
 456 impl fmt::Display for InstIdx {
 457     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 458         write!(f, "inst{}", self.get())
 459     }
 460 }
 461
 462 impl fmt::Display for BlockIdx {
 463     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 464         write!(f, "blk{}", self.get())
 465     }
 466 }
 467
 468 impl fmt::Display for BlockParamIdx {
 469     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 470         write!(f, "bparam{}", self.get())
 471     }
 472 }
 473
 474 impl fmt::Display for OperandIdx {
 475     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 476         write!(f, "operand{}", self.get())
 477     }
 478 }
 479
 480 impl fmt::Display for LiveRangeIdx {
 481     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 482         write!(f, "live_range{}", self.get())
 483     }
 484 }
 485
 486 impl BlockIdx {
 487     pub const ENTRY_BLOCK: BlockIdx = BlockIdx::new(0);
 488 }
 489
 490 /// range of instruction indexes from `start` inclusive to `end` exclusive.
 491 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 492 pub struct InstRange {
 493     pub start: InstIdx,
 494     pub end: InstIdx,
 495 }
 496
 497 impl fmt::Display for InstRange {
 498     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 499         write!(f, "inst{}..{}", self.start.get(), self.end.get())
 500     }
 501 }
 502
 503 impl IntoIterator for InstRange {
 504     type Item = InstIdx;
 505     type IntoIter = InstRangeIter;
 506     fn into_iter(self) -> Self::IntoIter {
 507         self.iter()
 508     }
 509 }
 510
 511 impl InstRange {
 512     pub const fn iter(self) -> InstRangeIter {
 513         InstRangeIter(self.start.get()..self.end.get())
 514     }
 515     pub const fn is_empty(self) -> bool {
 516         self.start.get() >= self.end.get()
 517     }
 518     pub const fn first(self) -> Option<InstIdx> {
 519         Some(const_try_opt!(self.split_first()).0)
 520     }
 521     pub const fn last(self) -> Option<InstIdx> {
 522         Some(const_try_opt!(self.split_last()).0)
 523     }
 524     pub const fn split_first(self) -> Option<(InstIdx, InstRange)> {
 525         if self.is_empty() {
 526             None
 527         } else {
 528             Some((
 529                 self.start,
 530                 Self {
 531                     start: self.start.next(),
 532                     end: self.end,
 533                 },
 534             ))
 535         }
 536     }
 537     pub const fn split_last(self) -> Option<(InstIdx, InstRange)> {
 538         if self.is_empty() {
 539             None
 540         } else {
 541             Some((
 542                 self.end.prev(),
 543                 Self {
 544                     start: self.start,
 545                     end: self.end.prev(),
 546                 },
 547             ))
 548         }
 549     }
 550     pub const fn len(self) -> usize {
 551         if self.is_empty() {
 552             0
 553         } else {
 554             self.end.get() - self.start.get()
 555         }
 556     }
 557 }
 558
 559 #[derive(Clone, Debug)]
 560 pub struct InstRangeIter(Range<usize>);
 561
 562 impl InstRangeIter {
 563     pub const fn range(self) -> InstRange {
 564         InstRange {
 565             start: InstIdx::new(self.0.start),
 566             end: InstIdx::new(self.0.end),
 567         }
 568     }
 569 }
 570
 571 impl Iterator for InstRangeIter {
 572     type Item = InstIdx;
 573
 574     fn next(&mut self) -> Option<Self::Item> {
 575         self.0.next().map(InstIdx::new)
 576     }
 577
 578     fn size_hint(&self) -> (usize, Option<usize>) {
 579         let v = self.0.len();
 580         (v, Some(v))
 581     }
 582
 583     fn count(self) -> usize
 584     where
 585         Self: Sized,
 586     {
 587         self.len()
 588     }
 589
 590     fn last(self) -> Option<Self::Item>
 591     where
 592         Self: Sized,
 593     {
 594         self.0.last().map(InstIdx::new)
 595     }
 596
 597     fn nth(&mut self, n: usize) -> Option<Self::Item> {
 598         self.0.nth(n).map(InstIdx::new)
 599     }
 600 }
 601
 602 impl FusedIterator for InstRangeIter {}
 603
 604 impl DoubleEndedIterator for InstRangeIter {
 605     fn next_back(&mut self) -> Option<Self::Item> {
 606         self.0.next_back().map(InstIdx::new)
 607     }
 608
 609     fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
 610         self.0.nth_back(n).map(InstIdx::new)
 611     }
 612 }
 613
 614 impl ExactSizeIterator for InstRangeIter {
 615     fn len(&self) -> usize {
 616         self.0.len()
 617     }
 618 }