From e3714ddd10a8864a7557baff858f38b2f3042474 Mon Sep 17 00:00:00 2001
From: Jakob Stoklund Olesen <jolesen@mozilla.com>
Date: Thu, 1 Feb 2018 17:10:21 -0800
Subject: [PATCH] Add a func.inst_offsets() iterator.

This Function method can be used after the final code layout has been
computed. It returns all the instructions in an EBB along with their
encoded size and offset from the beginning of the function.

This is useful for extracting additional metadata about trapping
instructions and other things that may be needed by a VM.
---
 cranelift/src/filetest/binemit.rs |  5 ++--
 lib/cretonne/src/ir/function.rs   | 46 ++++++++++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/cranelift/src/filetest/binemit.rs b/cranelift/src/filetest/binemit.rs
index aca308021b..9b18a439a1 100644
--- a/cranelift/src/filetest/binemit.rs
+++ b/cranelift/src/filetest/binemit.rs
@@ -211,7 +211,8 @@ impl SubTest for TestBinEmit {
                 "Inconsistent {} header offset",
                 ebb
             );
-            for inst in func.layout.ebb_insts(ebb) {
+            for (offset, inst, enc_bytes) in func.inst_offsets(ebb, &encinfo) {
+                assert_eq!(sink.offset, offset);
                 sink.text.clear();
                 let enc = func.encodings[inst];
 
@@ -234,7 +235,7 @@ impl SubTest for TestBinEmit {
                     // Verify the encoding recipe sizes against the ISAs emit_inst implementation.
                     assert_eq!(
                         emitted,
-                        encinfo.bytes(enc),
+                        enc_bytes,
                         "Inconsistent size for [{}] {}",
                         encinfo.display(enc),
                         func.dfg.display_inst(inst, isa)
diff --git a/lib/cretonne/src/ir/function.rs b/lib/cretonne/src/ir/function.rs
index 457b47a5e0..91ee7eca75 100644
--- a/lib/cretonne/src/ir/function.rs
+++ b/lib/cretonne/src/ir/function.rs
@@ -3,13 +3,14 @@
 //! The `Function` struct defined in this module owns all of its extended basic blocks and
 //! instructions.
 
+use binemit::CodeOffset;
 use entity::{PrimaryMap, EntityMap};
 use ir;
 use ir::{ExternalName, CallConv, Signature, DataFlowGraph, Layout};
 use ir::{InstEncodings, ValueLocations, JumpTables, StackSlots, EbbOffsets, SourceLocs};
 use ir::{Ebb, JumpTableData, JumpTable, StackSlotData, StackSlot, SigRef, ExtFuncData, FuncRef,
          GlobalVarData, GlobalVar, HeapData, Heap};
-use isa::TargetIsa;
+use isa::{TargetIsa, EncInfo};
 use std::fmt;
 use write::write_function;
 
@@ -153,6 +154,28 @@ impl Function {
             self.dfg.ebb_params(entry)[i]
         })
     }
+
+    /// Get an iterator over the instructions in `ebb`, including offsets and encoded instruction
+    /// sizes.
+    ///
+    /// The iterator returns `(offset, inst, size)` tuples, where `offset` if the offset in bytes
+    /// from the beginning of the function to the instruction, and `size` is the size of the
+    /// instruction in bytes, or 0 for unencoded instructions.
+    ///
+    /// This function can only be used after the code layout has been computed by the
+    /// `binemit::relax_branches()` function.
+    pub fn inst_offsets<'a>(&'a self, ebb: Ebb, encinfo: &EncInfo) -> InstOffsetIter<'a> {
+        assert!(
+            !self.offsets.is_empty(),
+            "Code layout must be computed first"
+        );
+        InstOffsetIter {
+            encinfo: encinfo.clone(),
+            encodings: &self.encodings,
+            offset: self.offsets[ebb],
+            iter: self.layout.ebb_insts(ebb),
+        }
+    }
 }
 
 /// Wrapper type capable of displaying a `Function` with correct ISA annotations.
@@ -175,3 +198,24 @@ impl fmt::Debug for Function {
         write_function(fmt, self, None)
     }
 }
+
+/// Iterator returning instruction offsets and sizes: `(offset, inst, size)`.
+pub struct InstOffsetIter<'a> {
+    encinfo: EncInfo,
+    encodings: &'a InstEncodings,
+    offset: CodeOffset,
+    iter: ir::layout::Insts<'a>,
+}
+
+impl<'a> Iterator for InstOffsetIter<'a> {
+    type Item = (CodeOffset, ir::Inst, CodeOffset);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next().map(|inst| {
+            let size = self.encinfo.bytes(self.encodings[inst]);
+            let offset = self.offset;
+            self.offset += size;
+            (offset, inst, size)
+        })
+    }
+}