[Mesa-dev] [PATCH] nir: fix nir_before/after_cf_node() for function impls

[Connor Abbott]

Signed-off-by: Connor Abbott <cwabbott0 at gmail.com>
---
 src/compiler/nir/nir.h | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 98451c6..8316b0d 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1995,19 +1995,33 @@ nir_after_block_before_jump(nir_block *block)
 static inline nir_cursor
 nir_before_cf_node(nir_cf_node *node)
 {
-   if (node->type == nir_cf_node_block)
+   switch(node->type) {
+   case nir_cf_node_block:
       return nir_before_block(nir_cf_node_as_block(node));
 
-   return nir_after_block(nir_cf_node_as_block(nir_cf_node_prev(node)));
+   case nir_cf_node_function:
+      return nir_before_block(
+         nir_start_block(nir_cf_node_as_function(node)));
+
+   default:
+      return nir_after_block(nir_cf_node_as_block(nir_cf_node_prev(node)));
+   }
 }
 
 static inline nir_cursor
 nir_after_cf_node(nir_cf_node *node)
 {
-   if (node->type == nir_cf_node_block)
+   switch(node->type) {
+   case nir_cf_node_block:
       return nir_after_block(nir_cf_node_as_block(node));
 
-   return nir_before_block(nir_cf_node_as_block(nir_cf_node_next(node)));
+   case nir_cf_node_function:
+      return nir_after_block(
+         nir_impl_last_block(nir_cf_node_as_function(node)));
+
+   default:
+      return nir_before_block(nir_cf_node_as_block(nir_cf_node_next(node)));
+   }
 }
 
 static inline nir_cursor
-- 
2.5.5