1 files changed, 31 insertions, 15 deletions

diff --git a/measure-stack.cc b/measure-stack.cc
index 5151772..e676735 100644
--- a/measure-stack.cc
+++ b/measure-stack.cc
@@ -67,8 +67,7 @@ immediate (Operand const& operand)
 }
 
 static int64_t
-stack_offset (Instruction const& instruction,
-              Architecture architecture)
+stack_offset (Instruction const& instruction, Architecture architecture)
 {
   switch (instruction.getOperation ().getID ())
     {
@@ -85,6 +84,21 @@ stack_offset (Instruction const& instruction,
     }
 }
 
+static size_t
+stack_offset (char const* buffer, size_t length,
+              RegisterAST::Ptr stack_pointer,
+              Architecture architecture)
+{
+  InstructionDecoder decoder {buffer, length, architecture};
+  size_t offset = 0;
+  for (auto instruction = decoder.decode ();
+       instruction.isValid ();
+       instruction = decoder.decode ())
+    if (instruction.isWritten (stack_pointer))
+      offset += stack_offset (instruction, architecture);
+  return offset;
+}
+
 int
 main (int argc, char** argv)
 {
@@ -108,21 +122,23 @@ main (int argc, char** argv)
       size_t stack_size = 0;
       for (auto* const fun : functions)
         {
-          auto const entry = fun->addr ();
-          auto const* buffer = (char*) cs.getPtrToInstruction (entry);
-          auto const length = fun->entry ()->end () - entry;
-          InstructionDecoder decoder {buffer, length, architecture};
-          size_t s = 0;
-          for (auto insn = decoder.decode ();
-               insn.isValid ();
-               insn = decoder.decode ())
-            if (insn.isWritten (stack_pointer))
-              s += stack_offset (insn, architecture);
-          if (s == 0)
+          size_t offset = 0;
+          for (auto* prologue = fun->entry ();
+               prologue;
+               prologue = (prologue->targets ().size () == 1
+                           ? (*prologue->targets ().begin ())->trg ()
+                           : nullptr))
+            {
+              auto const* const buffer = (char*)
+                cs.getPtrToInstruction (prologue->start ());
+              offset += stack_offset (buffer, prologue->size (),
+                                      stack_pointer, architecture);
+            }
+          if (offset == 0)
             continue;
           if (stack_size == 0)
-            stack_size = s;
-          else if (s != stack_size)
+            stack_size = offset;
+          else if (offset != stack_size)
             die_for (address, "functions with different stack sizes spanning");
         }
       std::cout << stack_size << '\n';