Give WalkStack's callback a way to terminate early.

Also combine WalkStack and WalkStackUntilUpcall.

Change-Id: Ida25665de72e5fd8e17946886a387b27cf841457
diff --git a/src/dalvik_system_VMStack.cc b/src/dalvik_system_VMStack.cc
index bcfe29b..a83ad45 100644
--- a/src/dalvik_system_VMStack.cc
+++ b/src/dalvik_system_VMStack.cc
@@ -59,18 +59,15 @@
   struct ClosestUserClassLoaderVisitor : public Thread::StackVisitor {
     ClosestUserClassLoaderVisitor(Object* bootstrap, Object* system)
       : bootstrap(bootstrap), system(system), class_loader(NULL) {}
-    virtual void VisitFrame(const Frame& f, uintptr_t) {
-      if (class_loader != NULL) {
-        // If we already found a result, nothing to do.
-        // TODO: need SmartFrame (Thread::WalkStack-like iterator).
-        // (or change VisitFrame to let us return bool to stop visiting)
-        return;
-      }
+    bool VisitFrame(const Frame& f, uintptr_t) {
+      DCHECK(class_loader == NULL);
       Class* c = f.GetMethod()->GetDeclaringClass();
       Object* cl = c->GetClassLoader();
       if (cl != NULL && cl != bootstrap && cl != system) {
         class_loader = cl;
+        return false;
       }
+      return true;
     }
     Object* bootstrap;
     Object* system;
diff --git a/src/debugger.cc b/src/debugger.cc
index 1480ee7..0255775 100644
--- a/src/debugger.cc
+++ b/src/debugger.cc
@@ -1468,11 +1468,11 @@
 static int GetStackDepth(Thread* thread) {
   struct CountStackDepthVisitor : public Thread::StackVisitor {
     CountStackDepthVisitor() : depth(0) {}
-    virtual void VisitFrame(const Frame& f, uintptr_t) {
-      // TODO: we'll need to skip callee-save frames too.
+    bool VisitFrame(const Frame& f, uintptr_t) {
       if (f.HasMethod()) {
         ++depth;
       }
+      return true;
     }
     size_t depth;
   };
@@ -1486,30 +1486,24 @@
   return GetStackDepth(DecodeThread(threadId));
 }
 
-bool Dbg::GetThreadFrame(JDWP::ObjectId threadId, int desired_frame_number, JDWP::FrameId* pFrameId, JDWP::JdwpLocation* pLoc) {
+void Dbg::GetThreadFrame(JDWP::ObjectId threadId, int desired_frame_number, JDWP::FrameId* pFrameId, JDWP::JdwpLocation* pLoc) {
   ScopedThreadListLock thread_list_lock;
   struct GetFrameVisitor : public Thread::StackVisitor {
     GetFrameVisitor(int desired_frame_number, JDWP::FrameId* pFrameId, JDWP::JdwpLocation* pLoc)
-        : found(false), depth(0), desired_frame_number(desired_frame_number), pFrameId(pFrameId), pLoc(pLoc) {
+        : depth(0), desired_frame_number(desired_frame_number), pFrameId(pFrameId), pLoc(pLoc) {
     }
-    virtual void VisitFrame(const Frame& f, uintptr_t pc) {
-      if (found) {
-        return;
-      }
-
-      // TODO: we'll need to skip callee-save frames too.
+    bool VisitFrame(const Frame& f, uintptr_t pc) {
       if (!f.HasMethod()) {
-        return; // The debugger can't do anything useful with a frame that has no Method*.
+        return true; // The debugger can't do anything useful with a frame that has no Method*.
       }
-
       if (depth == desired_frame_number) {
         *pFrameId = reinterpret_cast<JDWP::FrameId>(f.GetSP());
         SetLocation(*pLoc, f.GetMethod(), pc);
-        found = true;
+        return false;
       }
       ++depth;
+      return true;
     }
-    bool found;
     int depth;
     int desired_frame_number;
     JDWP::FrameId* pFrameId;
@@ -1518,7 +1512,6 @@
   GetFrameVisitor visitor(desired_frame_number, pFrameId, pLoc);
   visitor.desired_frame_number = desired_frame_number;
   DecodeThread(threadId)->WalkStack(&visitor);
-  return visitor.found;
 }
 
 JDWP::ObjectId Dbg::GetThreadSelfId() {
@@ -1940,8 +1933,7 @@
       gSingleStepControl.method = NULL;
       gSingleStepControl.stack_depth = 0;
     }
-    virtual void VisitFrame(const Frame& f, uintptr_t pc) {
-      // TODO: we'll need to skip callee-save frames too.
+    bool VisitFrame(const Frame& f, uintptr_t pc) {
       if (f.HasMethod()) {
         ++gSingleStepControl.stack_depth;
         if (gSingleStepControl.method == NULL) {
@@ -1955,6 +1947,7 @@
           }
         }
       }
+      return true;
     }
   };
   SingleStepStackVisitor visitor;
@@ -2840,17 +2833,16 @@
   explicit AllocRecordStackVisitor(AllocRecord* record) : record(record), depth(0) {
   }
 
-  virtual void VisitFrame(const Frame& f, uintptr_t pc) {
+  bool VisitFrame(const Frame& f, uintptr_t pc) {
     if (depth >= kMaxAllocRecordStackDepth) {
-      return;
+      return false;
     }
-    Method* m = f.GetMethod();
-    if (m == NULL || m->IsCalleeSaveMethod()) {
-      return;
+    if (f.HasMethod()) {
+      record->stack[depth].method = f.GetMethod();
+      record->stack[depth].raw_pc = pc;
+      ++depth;
     }
-    record->stack[depth].method = m;
-    record->stack[depth].raw_pc = pc;
-    ++depth;
+    return true;
   }
 
   ~AllocRecordStackVisitor() {
diff --git a/src/debugger.h b/src/debugger.h
index 659843a..fac9705 100644
--- a/src/debugger.h
+++ b/src/debugger.h
@@ -193,7 +193,7 @@
   static void GetThreadGroupThreads(JDWP::ObjectId threadGroupId, JDWP::ObjectId** ppThreadIds, uint32_t* pThreadCount);
   static void GetAllThreads(JDWP::ObjectId** ppThreadIds, uint32_t* pThreadCount);
   static int GetThreadFrameCount(JDWP::ObjectId threadId);
-  static bool GetThreadFrame(JDWP::ObjectId threadId, int num, JDWP::FrameId* pFrameId, JDWP::JdwpLocation* pLoc);
+  static void GetThreadFrame(JDWP::ObjectId threadId, int num, JDWP::FrameId* pFrameId, JDWP::JdwpLocation* pLoc);
 
   static JDWP::ObjectId GetThreadSelfId();
   static void SuspendVM();
diff --git a/src/jdwp/jdwp_handler.cc b/src/jdwp/jdwp_handler.cc
index 284ca64..9fb8ad0 100644
--- a/src/jdwp/jdwp_handler.cc
+++ b/src/jdwp/jdwp_handler.cc
@@ -976,6 +976,8 @@
   for (uint32_t i = start_frame; i < (start_frame + length); ++i) {
     FrameId frameId;
     JdwpLocation loc;
+    // TODO: switch to GetThreadFrames so we don't have to search for each frame
+    // even though we only want them in order.
     Dbg::GetThreadFrame(threadId, i, &frameId, &loc);
 
     expandBufAdd8BE(pReply, frameId);
diff --git a/src/thread.cc b/src/thread.cc
index 8470b8d..b511aa7 100644
--- a/src/thread.cc
+++ b/src/thread.cc
@@ -566,9 +566,9 @@
   virtual ~StackDumpVisitor() {
   }
 
-  void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     if (!frame.HasMethod()) {
-      return;
+      return true;
     }
     const int kMaxRepetition = 3;
     Method* m = frame.GetMethod();
@@ -606,6 +606,7 @@
     if (frame_count++ == 0) {
       Monitor::DescribeWait(os, thread);
     }
+    return true;
   }
   MethodHelper mh;
   Method* last_method;
@@ -1066,7 +1067,7 @@
  public:
   CountStackDepthVisitor() : depth_(0), skip_depth_(0), skipping_(true) {}
 
-  virtual void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     // We want to skip frames up to and including the exception's constructor.
     // Note we also skip the frame if it doesn't have a method (namely the callee
     // save frame)
@@ -1081,6 +1082,7 @@
     } else {
       ++skip_depth_;
     }
+    return true;
   }
 
   int GetDepth() const {
@@ -1127,20 +1129,21 @@
 
   virtual ~BuildInternalStackTraceVisitor() {}
 
-  virtual void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     if (method_trace_ == NULL || pc_trace_ == NULL) {
-      return; // We're probably trying to fillInStackTrace for an OutOfMemoryError.
+      return true; // We're probably trying to fillInStackTrace for an OutOfMemoryError.
     }
     if (skip_depth_ > 0) {
       skip_depth_--;
-      return;
+      return true;
     }
     if (!frame.HasMethod()) {
-      return;  // ignore callee save frames
+      return true;  // ignore callee save frames
     }
     method_trace_->Set(count_, frame.GetMethod());
     pc_trace_->Set(count_, pc);
     ++count_;
+    return true;
   }
 
   jobject GetInternalStackTrace() const {
@@ -1190,7 +1193,7 @@
   return sirt;
 }
 
-void Thread::WalkStack(StackVisitor* visitor) const {
+void Thread::WalkStack(StackVisitor* visitor, bool include_upcalls) const {
   Frame frame = GetTopOfStack();
   uintptr_t pc = ManglePc(top_of_managed_stack_pc_);
 #if defined(__arm__)
@@ -1200,10 +1203,13 @@
   // CHECK(native_to_managed_record_ != NULL);
   NativeToManagedRecord* record = native_to_managed_record_;
 
-  while (frame.GetSP() != 0) {
-    for ( ; frame.GetMethod() != 0; frame.Next()) {
+  while (frame.GetSP() != NULL) {
+    for ( ; frame.GetMethod() != NULL; frame.Next()) {
       // DCHECK(frame.GetMethod()->IsWithinCode(pc));  // TODO: restore IsWithinCode
-      visitor->VisitFrame(frame, pc);
+      bool should_continue = visitor->VisitFrame(frame, pc);
+      if (!should_continue) {
+        return;
+      }
       pc = ManglePc(frame.GetReturnPC());
       if (Runtime::Current()->IsMethodTracingActive()) {
 #if defined(__arm__)
@@ -1216,8 +1222,14 @@
 #endif
       }
     }
+    if (include_upcalls) {
+      bool should_continue = visitor->VisitFrame(frame, pc);
+      if (!should_continue) {
+        return;
+      }
+    }
     if (record == NULL) {
-      break;
+      return;
     }
     // last_tos should return Frame instead of sp?
     frame.SetSP(reinterpret_cast<Method**>(record->last_top_of_managed_stack_));
@@ -1226,22 +1238,6 @@
   }
 }
 
-void Thread::WalkStackUntilUpCall(StackVisitor* visitor, bool include_upcall) const {
-  Frame frame = GetTopOfStack();
-  uintptr_t pc = ManglePc(top_of_managed_stack_pc_);
-
-  if (frame.GetSP() != 0) {
-    for ( ; frame.GetMethod() != 0; frame.Next()) {
-      // DCHECK(frame.GetMethod()->IsWithinCode(pc));  // TODO: restore IsWithinCode
-      visitor->VisitFrame(frame, pc);
-      pc = ManglePc(frame.GetReturnPC());
-    }
-    if (include_upcall) {
-      visitor->VisitFrame(frame, pc);
-    }
-  }
-}
-
 jobject Thread::CreateInternalStackTrace(JNIEnv* env) const {
   // Compute depth of stack
   CountStackDepthVisitor count_visitor;
@@ -1430,59 +1426,54 @@
 class CatchBlockStackVisitor : public Thread::StackVisitor {
  public:
   CatchBlockStackVisitor(Class* to_find, Context* ljc)
-      : found_(false), to_find_(to_find), long_jump_context_(ljc), native_method_count_(0) {
+      : to_find_(to_find), long_jump_context_(ljc), native_method_count_(0) {
 #ifndef NDEBUG
     handler_pc_ = 0xEBADC0DE;
     handler_frame_.SetSP(reinterpret_cast<Method**>(0xEBADF00D));
 #endif
   }
 
-  virtual void VisitFrame(const Frame& fr, uintptr_t pc) {
-    if (!found_) {
-      Method* method = fr.GetMethod();
-      if (method == NULL) {
-        // This is the upcall, we remember the frame and last_pc so that we may
-        // long jump to them
-        handler_pc_ = DemanglePc(pc);
-        handler_frame_ = fr;
-        return;
-      }
-      uint32_t dex_pc = DexFile::kDexNoIndex;
-      if (method->IsCalleeSaveMethod()) {
-        // ignore callee save method
-      } else if (method->IsNative()) {
-        native_method_count_++;
-      } else {
-        // Unwind stack during method tracing
-        if (Runtime::Current()->IsMethodTracingActive()) {
+  bool VisitFrame(const Frame& fr, uintptr_t pc) {
+    Method* method = fr.GetMethod();
+    if (method == NULL) {
+      // This is the upcall, we remember the frame and last_pc so that we may
+      // long jump to them
+      handler_pc_ = DemanglePc(pc);
+      handler_frame_ = fr;
+      return false;
+    }
+    uint32_t dex_pc = DexFile::kDexNoIndex;
+    if (method->IsCalleeSaveMethod()) {
+      // ignore callee save method
+    } else if (method->IsNative()) {
+      native_method_count_++;
+    } else {
+      // Unwind stack during method tracing
+      if (Runtime::Current()->IsMethodTracingActive()) {
 #if defined(__arm__)
-          uintptr_t trace_exit = reinterpret_cast<uintptr_t>(art_trace_exit_from_code);
-          if (ManglePc(trace_exit) == pc) {
-            pc = ManglePc(TraceMethodUnwindFromCode(Thread::Current()));
-          }
+        uintptr_t trace_exit = reinterpret_cast<uintptr_t>(art_trace_exit_from_code);
+        if (ManglePc(trace_exit) == pc) {
+          pc = ManglePc(TraceMethodUnwindFromCode(Thread::Current()));
+        }
 #else
-          UNIMPLEMENTED(WARNING);
+        UNIMPLEMENTED(WARNING);
 #endif
-        }
-        dex_pc = method->ToDexPC(pc);
       }
-      if (dex_pc != DexFile::kDexNoIndex) {
-        uint32_t found_dex_pc = method->FindCatchBlock(to_find_, dex_pc);
-        if (found_dex_pc != DexFile::kDexNoIndex) {
-          found_ = true;
-          handler_pc_ = method->ToNativePC(found_dex_pc);
-          handler_frame_ = fr;
-        }
-      }
-      if (!found_) {
-        // Caller may be handler, fill in callee saves in context
-        long_jump_context_->FillCalleeSaves(fr);
+      dex_pc = method->ToDexPC(pc);
+    }
+    if (dex_pc != DexFile::kDexNoIndex) {
+      uint32_t found_dex_pc = method->FindCatchBlock(to_find_, dex_pc);
+      if (found_dex_pc != DexFile::kDexNoIndex) {
+        handler_pc_ = method->ToNativePC(found_dex_pc);
+        handler_frame_ = fr;
+        return false;
       }
     }
+    // Caller may be handler, fill in callee saves in context
+    long_jump_context_->FillCalleeSaves(fr);
+    return true;
   }
 
-  // Did we find a catch block yet?
-  bool found_;
   // The type of the exception catch block to find
   Class* to_find_;
   // Frame with found handler or last frame if no handler found
@@ -1511,7 +1502,7 @@
 
   Context* long_jump_context = GetLongJumpContext();
   CatchBlockStackVisitor catch_finder(exception->GetClass(), long_jump_context);
-  WalkStackUntilUpCall(&catch_finder, true);
+  WalkStack(&catch_finder, true);
 
   Method** sp;
   uintptr_t throw_native_pc;
@@ -1589,7 +1580,7 @@
     context_(context), root_visitor_(root_visitor), arg_(arg) {
   }
 
-  void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     Method* m = frame.GetMethod();
     if (false) {
       LOG(INFO) << "Visiting stack roots in " << PrettyMethod(m, false)
@@ -1641,6 +1632,7 @@
       }
     }
     context_->FillCalleeSaves(frame);
+    return true;
   }
 
  private:
diff --git a/src/thread.h b/src/thread.h
index 94659c4..51385f9 100644
--- a/src/thread.h
+++ b/src/thread.h
@@ -92,7 +92,8 @@
   class StackVisitor {
    public:
     virtual ~StackVisitor() {}
-    virtual void VisitFrame(const Frame& frame, uintptr_t pc) = 0;
+    // Return 'true' if we should continue to visit more frames, 'false' to stop.
+    virtual bool VisitFrame(const Frame& frame, uintptr_t pc) = 0;
   };
 
   // Creates a new thread.
@@ -399,7 +400,7 @@
     return ThreadOffset(OFFSETOF_MEMBER(Thread, top_sirt_));
   }
 
-  void WalkStack(StackVisitor* visitor) const;
+  void WalkStack(StackVisitor* visitor, bool include_upcalls = false) const;
 
   DebugInvokeReq* GetInvokeReq() {
     return debug_invoke_req_;
@@ -463,8 +464,6 @@
 
   static void ThreadExitCallback(void* arg);
 
-  void WalkStackUntilUpCall(StackVisitor* visitor, bool include_upcall) const;
-
   // Thin lock thread id. This is a small integer used by the thin lock implementation.
   // This is not to be confused with the native thread's tid, nor is it the value returned
   // by java.lang.Thread.getId --- this is a distinct value, used only for locking. One
diff --git a/test/ReferenceMap/stack_walk_refmap_jni.cc b/test/ReferenceMap/stack_walk_refmap_jni.cc
index b3f82db..f6d4ac3 100644
--- a/test/ReferenceMap/stack_walk_refmap_jni.cc
+++ b/test/ReferenceMap/stack_walk_refmap_jni.cc
@@ -43,10 +43,10 @@
   ReferenceMap2Visitor() {
   }
 
-  void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     Method* m = frame.GetMethod();
     if (!m || m->IsNative()) {
-      return;
+      return true;
     }
     LOG(INFO) << "At " << PrettyMethod(m, false);
 
@@ -54,11 +54,11 @@
 
     if (!pc) {
       // pc == NULL: m is either a native method or a phony method
-      return;
+      return true;
     }
     if (m->IsCalleeSaveMethod()) {
       LOG(WARNING) << "no PC for " << PrettyMethod(m);
-      return;
+      return true;
     }
 
     // Enable this to dump reference map to LOG(INFO)
@@ -156,6 +156,8 @@
       CHECK(ref_bitmap);
       CHECK_REGS_CONTAIN_REFS(8, 3, 2, 1, 0);  // v8: this, v3: y, v2: y, v1: x, v0: ex
     }
+
+    return true;
   }
 };
 
diff --git a/test/StackWalk/stack_walk_jni.cc b/test/StackWalk/stack_walk_jni.cc
index 69b4837..65b3284 100644
--- a/test/StackWalk/stack_walk_jni.cc
+++ b/test/StackWalk/stack_walk_jni.cc
@@ -42,7 +42,7 @@
   ReferenceMapVisitor() {
   }
 
-  void VisitFrame(const Frame& frame, uintptr_t pc) {
+  bool VisitFrame(const Frame& frame, uintptr_t pc) {
     Method* m = frame.GetMethod();
     CHECK(m != NULL);
     LOG(INFO) << "At " << PrettyMethod(m, false);
@@ -50,7 +50,7 @@
     if (m->IsCalleeSaveMethod() || m->IsNative()) {
       LOG(WARNING) << "no PC for " << PrettyMethod(m);
       CHECK_EQ(pc, 0u);
-      return;
+      return true;
     }
     verifier::PcToReferenceMap map(m->GetGcMap(), m->GetGcMapLength());
     const uint8_t* reg_bitmap = map.FindBitMap(m->ToDexPC(pc));
@@ -89,6 +89,8 @@
       }
     }
     LOG(INFO) << reinterpret_cast<const void*>(reg_bitmap);
+
+    return true;
   }
 };