Support implicit return from block_lambda_revival.

[cycript.git] / Replace.cpp

diff --git a/Replace.cpp b/Replace.cpp
index 74f04d9d489f7712d48366cd7837dc44ca35a32e..7b99022fd957473ffe52b8de2608284000483b5f 100644 (file)
--- a/Replace.cpp
+++ b/Replace.cpp
@@ -29,6 +29,11 @@ CYFunctionExpression *CYNonLocalize(CYContext &context, CYFunctionExpression *fu
     return function;
 }
 
+static void CYImplicitReturn(CYStatement *&code) {
+    if (CYStatement *&last = CYGetLast(code))
+        last = last->Return();
+}
+
 CYExpression *CYAdd::Replace(CYContext &context) {
     CYInfix::Replace(context);
 
@@ -95,6 +100,11 @@ CYExpression *CYAssignment::Replace(CYContext &context) {
     return this;
 }
 
+CYStatement *CYBlock::Return() {
+    CYImplicitReturn(code_);
+    return this;
+}
+
 CYStatement *CYBlock::Replace(CYContext &context) {
     context.ReplaceAll(code_);
     if (code_ == NULL)
@@ -305,6 +315,10 @@ CYExpression *CYEncodedType::Replace(CYContext &context) {
     return typed_->Replace(context);
 }
 
+CYStatement *CYExpress::Return() {
+    return $ CYReturn(expression_);
+}
+
 CYStatement *CYExpress::Replace(CYContext &context) {
     context.Replace(expression_);
     return this;
@@ -430,7 +444,8 @@ void CYFunction::Replace_(CYContext &context, bool outer) {
         Inject(context);
 
     CYThisScope *_this(context.this_);
-    context.this_ = CYGetLast(&this_);
+    context.this_ = &this_;
+    context.this_ = CYGetLast(context.this_);
 
     CYNonLocal *nonlocal(context.nonlocal_);
     CYNonLocal *nextlocal(context.nextlocal_);
@@ -453,6 +468,9 @@ void CYFunction::Replace_(CYContext &context, bool outer) {
     parameters_->Replace(context, code_);
     context.ReplaceAll(code_);
 
+    if (implicit_)
+        CYImplicitReturn(code_);
+
     if (CYIdentifier *identifier = this_.identifier_)
         code_ = $$->*
             $ CYVar($L1($ CYDeclaration(identifier, $ CYThis())))->*
@@ -501,6 +519,12 @@ CYIdentifier *CYIdentifier::Replace(CYContext &context) {
     return replace_;
 }
 
+CYStatement *CYIf::Return() {
+    CYImplicitReturn(true_);
+    CYImplicitReturn(false_);
+    return this;
+}
+
 CYStatement *CYIf::Replace(CYContext &context) {
     context.Replace(test_);
     context.ReplaceAll(true_);
@@ -709,12 +733,14 @@ CYStatement *CYReturn::Replace(CYContext &context) {
 }
 
 CYExpression *CYRubyBlock::Replace(CYContext &context) {
-    // XXX: this needs to do something much more epic to handle return
     return call_->AddArgument(context, proc_->Replace(context));
 }
 
 CYExpression *CYRubyProc::Replace(CYContext &context) {
-    return CYNonLocalize(context, $ CYFunctionExpression(NULL, parameters_, code_));
+    CYFunctionExpression *function($ CYFunctionExpression(NULL, parameters_, code_));
+    function = CYNonLocalize(context, function);
+    function->implicit_ = true;
+    return function;
 }
 
 CYScope::CYScope(bool transparent, CYContext &context) :
@@ -831,6 +857,10 @@ void CYScope::Close(CYContext &context, CYStatement *&statements) {
         }
 }
 
+CYStatement *CYStatement::Return() {
+    return this;
+}
+
 CYString *CYString::Concat(CYContext &context, CYString *rhs) const {
     size_t size(size_ + rhs->size_);
     char *value($ char[size + 1]);