Fix for the `ERROR: control reached end of function without RETURN` but now need to implement for the two other cases.