diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index a67cab08..b77ec2b3 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -634,10 +634,15 @@ static StackPopResult parser__reduce(Parser *self, StackVersion version,
   return pop;
 }
 
-static inline const TSParseAction *parser__reductions_after_sequence(
-  Parser *self, TSStateId start_state, const TreeArray *trees_below,
-  uint32_t tree_count_below, const TreeArray *trees_above,
-  TSSymbol lookahead_symbol, uint32_t *count) {
+static const TSParseAction *parser__reductions_after_sequence(
+  Parser *self,
+  TSStateId start_state,
+  const TreeArray *trees_below,
+  uint32_t tree_count_below,
+  const TreeArray *trees_above,
+  TSSymbol lookahead_symbol,
+  uint32_t *count
+) {
   TSStateId state = start_state;
   uint32_t child_count = 0;
   *count = 0;
@@ -687,11 +692,10 @@ static inline const TSParseAction *parser__reductions_after_sequence(
   return actions;
 }
 
-static StackIterateAction parser__repair_error_callback(
-  void *payload, TSStateId state, TreeArray *trees, uint32_t tree_count,
-  bool is_done, bool is_pending) {
-
-  ErrorRepairSession *session = (ErrorRepairSession *)payload;
+static StackIterateAction parser__repair_error_callback(void *payload, TSStateId state,
+                                                        const TreeArray *trees,
+                                                        uint32_t tree_count) {
+  ErrorRepairSession *session = payload;
   Parser *self = session->parser;
   TSSymbol lookahead_symbol = session->lookahead_symbol;
   ReduceActionSet *repairs = &self->reduce_actions;
@@ -957,8 +961,7 @@ static bool parser__do_potential_reductions(Parser *self, StackVersion version)
 }
 
 static StackIterateAction parser__skip_preceding_trees_callback(
-  void *payload, TSStateId state, TreeArray *trees, uint32_t tree_count,
-  bool is_done, bool is_pending) {
+  void *payload, TSStateId state, const TreeArray *trees, uint32_t tree_count) {
   if (tree_count > 0 && state != ERROR_STATE) {
     uint32_t bytes_skipped = 0;
     for (uint32_t i = 0; i < trees->size; i++) {
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 385a30d6..b28c9217 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -47,6 +47,11 @@ typedef struct {
   bool found_valid_path;
 } StackPopSession;
 
+typedef struct {
+  void *payload;
+  StackIterateCallback callback;
+} StackIterateSession;
+
 typedef Array(StackNode *) StackNodeArray;
 
 typedef struct {
@@ -65,6 +70,8 @@ struct Stack {
   StackNode *base_node;
 };
 
+typedef StackIterateAction (*StackIterateInternalCallback)(void *, const Iterator *);
+
 static void stack_node_retain(StackNode *self) {
   if (!self)
     return;
@@ -214,7 +221,7 @@ static void ts_stack__add_slice(Stack *self, StackNode *node, TreeArray *trees,
 }
 
 INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
-                                  StackIterateCallback callback, void *payload) {
+                                  StackIterateInternalCallback callback, void *payload) {
   array_clear(&self->slices);
   array_clear(&self->iterators);
 
@@ -235,17 +242,8 @@ INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
     for (uint32_t i = 0, size = self->iterators.size; i < size; i++) {
       Iterator *iterator = &self->iterators.contents[i];
       StackNode *node = iterator->node;
-      bool is_done = node == self->base_node;
-
-      StackIterateAction action = callback(
-        payload,
-        node->state,
-        &iterator->trees,
-        iterator->tree_count,
-        is_done,
-        iterator->is_pending
-      );
 
+      StackIterateAction action = callback(payload, iterator);
       bool should_pop = action & StackIteratePop;
       bool should_stop = action & StackIterateStop || node->link_count == 0;
 
@@ -411,22 +409,26 @@ void ts_stack_push(Stack *self, StackVersion version, Tree *tree,
   head->node = new_node;
 }
 
-StackPopResult ts_stack_iterate(Stack *self, StackVersion version,
-                                StackIterateCallback callback, void *payload) {
-  return stack__iter(self, version, callback, payload);
+INLINE StackIterateAction iterate_callback(void *payload, const Iterator *iterator) {
+  StackIterateSession *session = payload;
+  return session->callback(session->payload, iterator->node->state, &iterator->trees, iterator->tree_count);
 }
 
-INLINE StackIterateAction pop_count_callback(void *payload, TSStateId state,
-                                             TreeArray *trees, uint32_t tree_count,
-                                             bool is_done, bool is_pending) {
+StackPopResult ts_stack_iterate(Stack *self, StackVersion version,
+                                StackIterateCallback callback, void *payload) {
+  StackIterateSession session = {payload, callback};
+  return stack__iter(self, version, iterate_callback, &session);
+}
+
+INLINE StackIterateAction pop_count_callback(void *payload, const Iterator *iterator) {
   StackPopSession *pop_session = (StackPopSession *)payload;
 
-  if (tree_count == pop_session->goal_tree_count) {
+  if (iterator->tree_count == pop_session->goal_tree_count) {
     pop_session->found_valid_path = true;
     return StackIteratePop | StackIterateStop;
   }
 
-  if (state == ERROR_STATE) {
+  if (iterator->node->state == ERROR_STATE) {
     if (pop_session->found_valid_path || pop_session->found_error) {
       return StackIterateStop;
     } else {
@@ -466,12 +468,9 @@ StackPopResult ts_stack_pop_count(Stack *self, StackVersion version,
   return pop;
 }
 
-INLINE StackIterateAction pop_pending_callback(void *payload, TSStateId state,
-                                               TreeArray *trees,
-                                               uint32_t tree_count, bool is_done,
-                                               bool is_pending) {
-  if (tree_count >= 1) {
-    if (is_pending) {
+INLINE StackIterateAction pop_pending_callback(void *payload, const Iterator *iterator) {
+  if (iterator->tree_count >= 1) {
+    if (iterator->is_pending) {
       return StackIteratePop | StackIterateStop;
     } else {
       return StackIterateStop;
@@ -490,10 +489,10 @@ StackPopResult ts_stack_pop_pending(Stack *self, StackVersion version) {
   return pop;
 }
 
-INLINE StackIterateAction pop_all_callback(void *payload, TSStateId state,
-                                           TreeArray *trees, uint32_t tree_count,
-                                           bool is_done, bool is_pending) {
-  return is_done ? (StackIteratePop | StackIterateStop) : StackIterateNone;
+INLINE StackIterateAction pop_all_callback(void *payload, const Iterator *iterator) {
+  return iterator->node->link_count == 0 ?
+    (StackIteratePop | StackIterateStop) :
+    StackIterateNone;
 }
 
 StackPopResult ts_stack_pop_all(Stack *self, StackVersion version) {
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index f372fe41..1eeb5cd0 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -12,8 +12,7 @@ extern "C" {
 
 typedef struct Stack Stack;
 
-typedef unsigned int StackVersion;
-
+typedef unsigned StackVersion;
 #define STACK_VERSION_NONE ((StackVersion)-1)
 
 typedef struct {
@@ -28,19 +27,16 @@ typedef struct {
   StackSliceArray slices;
 } StackPopResult;
 
+typedef unsigned StackIterateAction;
 enum {
   StackIterateNone,
-  StackIterateStop = 1 << 0,
-  StackIteratePop = 1 << 1,
+  StackIterateStop = 1,
+  StackIteratePop = 2,
 };
 
-typedef unsigned int StackIterateAction;
-
 typedef StackIterateAction (*StackIterateCallback)(void *, TSStateId state,
-                                                   TreeArray *trees,
-                                                   uint32_t tree_count,
-                                                   bool is_done,
-                                                   bool is_pending);
+                                                   const TreeArray *trees,
+                                                   uint32_t tree_count);
 
 /*
  *  Create a parse stack.
@@ -91,8 +87,7 @@ void ts_stack_push(Stack *, StackVersion, Tree *, bool, TSStateId);
  */
 StackPopResult ts_stack_pop_count(Stack *, StackVersion, uint32_t count);
 
-StackPopResult ts_stack_iterate(Stack *, StackVersion, StackIterateCallback,
-                                void *);
+StackPopResult ts_stack_iterate(Stack *, StackVersion, StackIterateCallback, void *);
 
 StackPopResult ts_stack_pop_pending(Stack *, StackVersion);
 
diff --git a/test/runtime/stack_test.cc b/test/runtime/stack_test.cc
index ecf579e0..03bf7091 100644
--- a/test/runtime/stack_test.cc
+++ b/test/runtime/stack_test.cc
@@ -54,11 +54,12 @@ vector<StackEntry> get_stack_entries(Stack *stack, StackVersion version) {
   ts_stack_iterate(
     stack,
     version,
-    [](void *payload, TSStateId state, TreeArray *trees, uint32_t tree_count, bool is_done, bool is_pending) -> StackIterateAction {
+    [](void *payload, TSStateId state, const TreeArray *trees, uint32_t tree_count) -> StackIterateAction {
       auto entries = static_cast<vector<StackEntry> *>(payload);
       StackEntry entry = {state, tree_count};
-      if (find(entries->begin(), entries->end(), entry) == entries->end())
+      if (find(entries->begin(), entries->end(), entry) == entries->end()) {
         entries->push_back(entry);
+      }
       return StackIterateNone;
     }, &result);
   return result;