diff --git a/fs/btrfs/extent-io-tree.c b/fs/btrfs/extent-io-tree.c
index 32f86a524bdd..f39eb5e816eb 100644
--- a/fs/btrfs/extent-io-tree.c
+++ b/fs/btrfs/extent-io-tree.c
@@ -283,9 +283,10 @@ static inline struct rb_node *tree_search_for_insert(struct extent_io_tree *tree
 /*
  * Inexact rb-tree search, return the next entry if @offset is not found
  */
-static inline struct rb_node *tree_search(struct extent_io_tree *tree, u64 offset)
+static inline struct extent_state *tree_search(struct extent_io_tree *tree, u64 offset)
 {
-	return tree_search_for_insert(tree, offset, NULL, NULL);
+	struct rb_node *n = tree_search_for_insert(tree, offset, NULL, NULL);
+	return (n) ? rb_entry(n, struct extent_state, rb_node) : NULL;
 }
 
 /**
@@ -591,7 +592,6 @@ int __clear_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 	struct extent_state *state;
 	struct extent_state *cached;
 	struct extent_state *prealloc = NULL;
-	struct rb_node *node;
 	u64 last_end;
 	int err;
 	int clear = 0;
@@ -642,10 +642,9 @@ int __clear_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 	 * this search will find the extents that end after
 	 * our range starts
 	 */
-	node = tree_search(tree, start);
-	if (!node)
+	state = tree_search(tree, start);
+	if (!state)
 		goto out;
-	state = rb_entry(node, struct extent_state, rb_node);
 hit_next:
 	if (state->start > end)
 		goto out;
@@ -758,7 +757,6 @@ static void wait_on_state(struct extent_io_tree *tree,
 void wait_extent_bit(struct extent_io_tree *tree, u64 start, u64 end, u32 bits)
 {
 	struct extent_state *state;
-	struct rb_node *node;
 
 	btrfs_debug_check_extent_io_range(tree, start, end);
 
@@ -769,11 +767,10 @@ void wait_extent_bit(struct extent_io_tree *tree, u64 start, u64 end, u32 bits)
 		 * this search will find all the extents that end after
 		 * our range starts
 		 */
-		node = tree_search(tree, start);
-		if (!node)
-			break;
-		state = rb_entry(node, struct extent_state, rb_node);
+		state = tree_search(tree, start);
 process_node:
+		if (!state)
+			break;
 		if (state->start > end)
 			goto out;
 
@@ -1313,25 +1310,18 @@ int lock_extent_bits(struct extent_io_tree *tree, u64 start, u64 end,
 static struct extent_state *
 find_first_extent_bit_state(struct extent_io_tree *tree, u64 start, u32 bits)
 {
-	struct rb_node *node;
 	struct extent_state *state;
 
 	/*
 	 * this search will find all the extents that end after
 	 * our range starts.
 	 */
-	node = tree_search(tree, start);
-	if (!node)
-		goto out;
-	state = rb_entry(node, struct extent_state, rb_node);
-
+	state = tree_search(tree, start);
 	while (state) {
-		state = rb_entry(node, struct extent_state, rb_node);
 		if (state->end >= start && (state->state & bits))
 			return state;
 		state = next_state(state);
 	}
-out:
 	return NULL;
 }
 
@@ -1541,7 +1531,6 @@ bool btrfs_find_delalloc_range(struct extent_io_tree *tree, u64 *start,
 			       u64 *end, u64 max_bytes,
 			       struct extent_state **cached_state)
 {
-	struct rb_node *node;
 	struct extent_state *state;
 	u64 cur_start = *start;
 	bool found = false;
@@ -1553,13 +1542,12 @@ bool btrfs_find_delalloc_range(struct extent_io_tree *tree, u64 *start,
 	 * this search will find all the extents that end after
 	 * our range starts.
 	 */
-	node = tree_search(tree, cur_start);
-	if (!node) {
+	state = tree_search(tree, cur_start);
+	if (!state) {
 		*end = (u64)-1;
 		goto out;
 	}
 
-	state = rb_entry(node, struct extent_state, rb_node);
 	while (state) {
 		if (found && (state->start != cur_start ||
 			      (state->state & EXTENT_BOUNDARY))) {
@@ -1597,7 +1585,6 @@ u64 count_range_bits(struct extent_io_tree *tree,
 		     u64 *start, u64 search_end, u64 max_bytes,
 		     u32 bits, int contig)
 {
-	struct rb_node *node;
 	struct extent_state *state;
 	u64 cur_start = *start;
 	u64 total_bytes = 0;
@@ -1616,11 +1603,7 @@ u64 count_range_bits(struct extent_io_tree *tree,
 	 * this search will find all the extents that end after
 	 * our range starts.
 	 */
-	node = tree_search(tree, cur_start);
-	if (!node)
-		goto out;
-
-	state = rb_entry(node, struct extent_state, rb_node);
+	state = tree_search(tree, cur_start);
 	while (state) {
 		if (state->start > search_end)
 			break;
@@ -1656,19 +1639,14 @@ int test_range_bit(struct extent_io_tree *tree, u64 start, u64 end,
 		   u32 bits, int filled, struct extent_state *cached)
 {
 	struct extent_state *state = NULL;
-	struct rb_node *node;
 	int bitset = 0;
 
 	spin_lock(&tree->lock);
 	if (cached && extent_state_in_tree(cached) && cached->start <= start &&
 	    cached->end > start)
-		node = &cached->rb_node;
+		state = cached;
 	else
-		node = tree_search(tree, start);
-	if (!node)
-		goto out;
-
-	state = rb_entry(node, struct extent_state, rb_node);
+		state = tree_search(tree, start);
 	while (state && start <= end) {
 		if (filled && state->start > start) {
 			bitset = 0;
@@ -1699,7 +1677,6 @@ int test_range_bit(struct extent_io_tree *tree, u64 start, u64 end,
 	/* We ran out of states and were still inside of our range. */
 	if (filled && !state)
 		bitset = 0;
-out:
 	spin_unlock(&tree->lock);
 	return bitset;
 }
