diff --git a/fs/btrfs/extent-io-tree.c b/fs/btrfs/extent-io-tree.c
index f39eb5e816eb..82a18c233fa0 100644
--- a/fs/btrfs/extent-io-tree.c
+++ b/fs/btrfs/extent-io-tree.c
@@ -227,6 +227,15 @@ int try_lock_extent(struct extent_io_tree *tree, u64 start, u64 end)
 	return 1;
 }
 
+static struct extent_state *next_state(struct extent_state *state)
+{
+	struct rb_node *next = rb_next(&state->rb_node);
+	if (next)
+		return rb_entry(next, struct extent_state, rb_node);
+	else
+		return NULL;
+}
+
 /**
  * Search @tree for an entry that contains @offset. Such entry would have
  * entry->start <= offset && entry->end >= offset.
@@ -244,15 +253,15 @@ int try_lock_extent(struct extent_io_tree *tree, u64 start, u64 end)
  * If no such entry exists, return pointer to entry that ends before @offset
  * and fill parameters @node_ret and @parent_ret, ie. does not return NULL.
  */
-static inline struct rb_node *tree_search_for_insert(struct extent_io_tree *tree,
-					             u64 offset,
-						     struct rb_node ***node_ret,
-						     struct rb_node **parent_ret)
+static inline struct extent_state *tree_search_for_insert(struct extent_io_tree *tree,
+							  u64 offset,
+							  struct rb_node ***node_ret,
+							  struct rb_node **parent_ret)
 {
 	struct rb_root *root = &tree->state;
 	struct rb_node **node = &root->rb_node;
 	struct rb_node *prev = NULL;
-	struct extent_state *entry;
+	struct extent_state *entry = NULL;
 
 	while (*node) {
 		prev = *node;
@@ -263,7 +272,7 @@ static inline struct rb_node *tree_search_for_insert(struct extent_io_tree *tree
 		else if (offset > entry->end)
 			node = &(*node)->rb_right;
 		else
-			return *node;
+			return entry;
 	}
 
 	if (node_ret)
@@ -272,12 +281,10 @@ static inline struct rb_node *tree_search_for_insert(struct extent_io_tree *tree
 		*parent_ret = prev;
 
 	/* Search neighbors until we find the first one past the end */
-	while (prev && offset > entry->end) {
-		prev = rb_next(prev);
-		entry = rb_entry(prev, struct extent_state, rb_node);
-	}
+	while (entry && offset > entry->end)
+		entry = next_state(entry);
 
-	return prev;
+	return entry;
 }
 
 /*
@@ -285,8 +292,7 @@ static inline struct rb_node *tree_search_for_insert(struct extent_io_tree *tree
  */
 static inline struct extent_state *tree_search(struct extent_io_tree *tree, u64 offset)
 {
-	struct rb_node *n = tree_search_for_insert(tree, offset, NULL, NULL);
-	return (n) ? rb_entry(n, struct extent_state, rb_node) : NULL;
+	return tree_search_for_insert(tree, offset, NULL, NULL);
 }
 
 /**
@@ -511,15 +517,6 @@ static int split_state(struct extent_io_tree *tree, struct extent_state *orig,
 	return 0;
 }
 
-static struct extent_state *next_state(struct extent_state *state)
-{
-	struct rb_node *next = rb_next(&state->rb_node);
-	if (next)
-		return rb_entry(next, struct extent_state, rb_node);
-	else
-		return NULL;
-}
-
 /*
  * utility function to clear some bits in an extent state struct.
  * it will optionally wake up anyone waiting on this state (wake == 1).
@@ -850,7 +847,6 @@ int set_extent_bit(struct extent_io_tree *tree, u64 start, u64 end, u32 bits,
 {
 	struct extent_state *state;
 	struct extent_state *prealloc = NULL;
-	struct rb_node *node;
 	struct rb_node **p;
 	struct rb_node *parent;
 	int err = 0;
@@ -880,17 +876,15 @@ int set_extent_bit(struct extent_io_tree *tree, u64 start, u64 end, u32 bits,
 	if (cached_state && *cached_state) {
 		state = *cached_state;
 		if (state->start <= start && state->end > start &&
-		    extent_state_in_tree(state)) {
-			node = &state->rb_node;
+		    extent_state_in_tree(state))
 			goto hit_next;
-		}
 	}
 	/*
 	 * this search will find all the extents that end after
 	 * our range starts.
 	 */
-	node = tree_search_for_insert(tree, start, &p, &parent);
-	if (!node) {
+	state = tree_search_for_insert(tree, start, &p, &parent);
+	if (!state) {
 		prealloc = alloc_extent_state_atomic(prealloc);
 		BUG_ON(!prealloc);
 		prealloc->start = start;
@@ -900,7 +894,6 @@ int set_extent_bit(struct extent_io_tree *tree, u64 start, u64 end, u32 bits,
 		prealloc = NULL;
 		goto out;
 	}
-	state = rb_entry(node, struct extent_state, rb_node);
 hit_next:
 	last_start = state->start;
 	last_end = state->end;
@@ -1086,7 +1079,6 @@ int convert_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 {
 	struct extent_state *state;
 	struct extent_state *prealloc = NULL;
-	struct rb_node *node;
 	struct rb_node **p;
 	struct rb_node *parent;
 	int err = 0;
@@ -1116,18 +1108,16 @@ int convert_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 	if (cached_state && *cached_state) {
 		state = *cached_state;
 		if (state->start <= start && state->end > start &&
-		    extent_state_in_tree(state)) {
-			node = &state->rb_node;
+		    extent_state_in_tree(state))
 			goto hit_next;
-		}
 	}
 
 	/*
 	 * this search will find all the extents that end after
 	 * our range starts.
 	 */
-	node = tree_search_for_insert(tree, start, &p, &parent);
-	if (!node) {
+	state = tree_search_for_insert(tree, start, &p, &parent);
+	if (!state) {
 		prealloc = alloc_extent_state_atomic(prealloc);
 		if (!prealloc) {
 			err = -ENOMEM;
@@ -1140,7 +1130,6 @@ int convert_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 		prealloc = NULL;
 		goto out;
 	}
-	state = rb_entry(node, struct extent_state, rb_node);
 hit_next:
 	last_start = state->start;
 	last_end = state->end;
