From 5fe9429a633869794dc6a38b4d39ff4693eb02ac Mon Sep 17 00:00:00 2001
From: Dzmitry Malyshau <kvarkus@gmail.com>
Date: Fri, 4 Dec 2020 15:03:53 -0500
Subject: [PATCH] Terminator processor, used in WGSL and GLSL frontends

---
 src/front/glsl/parser.rs | 19 ++++++++++---------
 src/front/wgsl/mod.rs    |  4 +++-
 src/proc/mod.rs          |  2 ++
 src/proc/terminator.rs   | 40 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 10 deletions(-)
 create mode 100644 src/proc/terminator.rs

diff --git a/src/front/glsl/parser.rs b/src/front/glsl/parser.rs
index db4935b0b..29cbbd3fc 100644
--- a/src/front/glsl/parser.rs
+++ b/src/front/glsl/parser.rs
@@ -5,10 +5,15 @@ pomelo! {
     //%verbose;
     %include {
         use super::super::{error::ErrorKind, token::*, ast::*};
-        use crate::{proc::Typifier, Arena, BinaryOperator, Binding, Block, Constant,
-            ConstantInner, EntryPoint, Expression, FallThrough, FastHashMap, Function, GlobalVariable, Handle, Interpolation,
-            LocalVariable, MemberOrigin, SampleLevel, ScalarKind, Statement, StorageAccess,
-            StorageClass, StructMember, Type, TypeInner, UnaryOperator};
+        use crate::{
+            proc::{ensure_block_returns, Typifier},
+            Arena, BinaryOperator, Binding, Block, Constant,
+            ConstantInner, EntryPoint, Expression, FallThrough,
+            FastHashMap, Function, GlobalVariable, Handle, Interpolation,
+            LocalVariable, MemberOrigin, SampleLevel, ScalarKind,
+            Statement, StorageAccess,
+            StorageClass, StructMember, Type, TypeInner, UnaryOperator,
+        };
     }
     %token #[derive(Debug)] #[cfg_attr(test, derive(PartialEq))] pub enum Token {};
     %parser pub struct Parser<'a> {};
@@ -1117,11 +1122,7 @@ pomelo! {
         extra.context.clear_scopes();
         extra.context.lookup_global_var_exps.clear();
         extra.context.typifier = Typifier::new();
-        // make sure function ends with return
-        match cs.last() {
-            Some(Statement::Return {..}) => {}
-            _ => {cs.push(Statement::Return { value:None });}
-        }
+        ensure_block_returns(&mut cs);
         f.body = cs;
         f.fill_global_use(&extra.module.global_variables);
         f
diff --git a/src/front/wgsl/mod.rs b/src/front/wgsl/mod.rs
index 08c2d1c3f..e2c95ee77 100644
--- a/src/front/wgsl/mod.rs
+++ b/src/front/wgsl/mod.rs
@@ -7,7 +7,7 @@ mod lexer;
 
 use crate::{
     arena::{Arena, Handle},
-    proc::{ResolveContext, ResolveError, Typifier},
+    proc::{ensure_block_returns, ResolveContext, ResolveError, Typifier},
     FastHashMap,
 };
 
@@ -1594,6 +1594,8 @@ impl Parser {
                 arguments: &fun.arguments,
             },
         )?;
+        // fixup the IR
+        ensure_block_returns(&mut fun.body);
         // done
         fun.fill_global_use(&module.global_variables);
         self.scopes.pop();
diff --git a/src/proc/mod.rs b/src/proc/mod.rs
index 961e55da7..d8388b4c5 100644
--- a/src/proc/mod.rs
+++ b/src/proc/mod.rs
@@ -4,6 +4,7 @@
 mod call_graph;
 mod interface;
 mod namer;
+mod terminator;
 mod typifier;
 mod validator;
 
@@ -11,6 +12,7 @@ mod validator;
 pub use call_graph::{CallGraph, CallGraphBuilder};
 pub use interface::{Interface, Visitor};
 pub use namer::{EntryPointIndex, NameKey, Namer};
+pub use terminator::ensure_block_returns;
 pub use typifier::{check_constant_type, ResolveContext, ResolveError, Typifier};
 pub use validator::{ValidationError, Validator};
 
diff --git a/src/proc/terminator.rs b/src/proc/terminator.rs
new file mode 100644
index 000000000..8b4d6ff33
--- /dev/null
+++ b/src/proc/terminator.rs
@@ -0,0 +1,40 @@
+/// Ensure that the given block has return statements
+/// at the end of its control flow.
+///
+/// Note: we don't want to blindly append a return statement
+/// to the end, because it may be either redundant or invalid,
+/// e.g. when the user already has returns in if/else branches.
+pub fn ensure_block_returns(block: &mut crate::Block) {
+    match block.last_mut() {
+        Some(&mut crate::Statement::Block(ref mut b)) => {
+            ensure_block_returns(b);
+        }
+        Some(&mut crate::Statement::If {
+            condition: _,
+            ref mut accept,
+            ref mut reject,
+        }) => {
+            ensure_block_returns(accept);
+            ensure_block_returns(reject);
+        }
+        Some(&mut crate::Statement::Switch {
+            selector: _,
+            ref mut cases,
+            ref mut default,
+        }) => {
+            for case in cases.values_mut() {
+                if let (ref mut b, None) = *case {
+                    ensure_block_returns(b);
+                }
+            }
+            ensure_block_returns(default);
+        }
+        Some(&mut crate::Statement::Break)
+        | Some(&mut crate::Statement::Continue)
+        | Some(&mut crate::Statement::Return { .. })
+        | Some(&mut crate::Statement::Kill) => (),
+        Some(&mut crate::Statement::Loop { .. })
+        | Some(&mut crate::Statement::Store { .. })
+        | None => block.push(crate::Statement::Return { value: None }),
+    }
+}