From 71f19068a7b13eec860824c9bbb0bcf38a2ee769 Mon Sep 17 00:00:00 2001
From: Alpha Chen <alpha@kejadlen.dev>
Date: Thu, 18 Aug 2022 19:20:14 -0700
Subject: [PATCH] 11.5.1

---
 ruby/lib/lox/resolver.rb       | 27 ++++++++++++++++++++-------
 ruby/test/lox/test_resolver.rb |  8 ++++++++
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/ruby/lib/lox/resolver.rb b/ruby/lib/lox/resolver.rb
index 19520d6..ac5aa27 100644
--- a/ruby/lib/lox/resolver.rb
+++ b/ruby/lib/lox/resolver.rb
@@ -7,6 +7,7 @@ module Lox
       @interpreter = interpreter
 
       @scopes = []
+      @current_func = :NONE
     end
 
     def resolve(*resolvees)
@@ -29,7 +30,7 @@ module Lox
       declare(stmt.name)
       define(stmt.name)
 
-      resolve_function(stmt)
+      resolve_function(stmt, :FUNCTION)
       nil
     end
 
@@ -40,6 +41,7 @@ module Lox
 
     def visit_print(stmt) = resolve(stmt.expr)
     def visit_return(stmt)
+      raise ResolverError.new(stmt.keyword, "Can't return from top-level code.") if @current_func == :NONE
       resolve(stmt.value) if stmt.value
     end
 
@@ -106,15 +108,26 @@ module Lox
       @interpreter.resolve(expr, depth)
     end
 
-    def resolve_function(fn)
-      with_scope do
-        fn.params.each do |param|
-          declare(param)
-          define(param)
+    def resolve_function(fn, type)
+      with_func_type(type) do
+        with_scope do
+          fn.params.each do |param|
+            declare(param)
+            define(param)
+          end
+          resolve(*fn.body)
         end
-        resolve(*fn.body)
       end
     end
 
+    def with_func_type(type)
+      enclosing_func = @current_func
+      @current_func = type
+
+      yield
+
+      @current_func = enclosing_func
+    end
+
   end
 end
diff --git a/ruby/test/lox/test_resolver.rb b/ruby/test/lox/test_resolver.rb
index 5509e86..81706eb 100644
--- a/ruby/test/lox/test_resolver.rb
+++ b/ruby/test/lox/test_resolver.rb
@@ -51,6 +51,14 @@ class TestResolver < Lox::Test
     end
   end
 
+  def test_returning_from_top_level
+    assert_raises Lox::ResolverError do
+      resolve(<<~SRC)
+        return;
+      SRC
+    end
+  end
+
   private
 
   def assert_resolved(expected, src)