fix main return bug (#6)

* fix main return bug * fix test case * fix options_test * fix operator_test
Vilsol · Apr 7, 2021 · 9038703 · 9038703
1 parent 465fc39
commit 9038703
Show file tree

Hide file tree

Showing 16 changed files with 44 additions and 6 deletions.
diff --git a/tests/base_test.go b/tests/base_test.go
@@ -64,7 +64,7 @@ func TestBase(t *testing.T) {
 				t.Error(err)
 				return
 			}
-
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/constant_test.go b/tests/constant_test.go
@@ -38,6 +38,7 @@ print x`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/control_test.go b/tests/control_test.go
@@ -45,6 +45,7 @@ func TestControl(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/custom_test.go b/tests/custom_test.go
@@ -30,6 +30,7 @@ func TestCustom(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/draw_test.go b/tests/draw_test.go
@@ -80,6 +80,7 @@ func TestDraw(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/extra_test.go b/tests/extra_test.go
@@ -31,6 +31,7 @@ jump 1 lessThan @time _main_0`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/native_test.go b/tests/native_test.go
@@ -45,6 +45,7 @@ print "\n"`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/operator_test.go b/tests/operator_test.go
@@ -73,6 +73,7 @@ print 1`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}
@@ -201,6 +202,7 @@ op not _main_x _main_0 -1`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}
@@ -244,6 +246,7 @@ func TestFunctionOperator(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/options_test.go b/tests/options_test.go
@@ -47,7 +47,8 @@ set _main_0 @return_0      	# 11
 print _main_0              	# 12	
 print "\n"                 	# 13	
 op add _main_i _main_i 1   	# 14	
-jump 8 lessThan _main_i 10 	# 15	`,
+jump 8 lessThan _main_i 10 	# 15	
+end                        	# 16	`,
 		},
 		{
 			name:  "Comments",
@@ -77,7 +78,8 @@ set _main_0 @return_0      	# Set variable to returned value
 print _main_0              	# Call to native function       	
 print "\n"                 	# Call to native function       	
 op add _main_i _main_i 1   	# Execute increment/decrement   	
-jump 8 lessThan _main_i 10 	# Jump to start of loop         	`,
+jump 8 lessThan _main_i 10 	# Jump to start of loop         	
+end                        	# Trampoline back               	`,
 		},
 		{
 			name:  "Comments",
@@ -101,7 +103,8 @@ set _main_0 @return_0
 print _main_0              	# println(foo(i))	
 print "\n"                 	# println(foo(i))	
 op add _main_i _main_i 1   	# i++            	
-jump 8 lessThan _main_i 10 	                 	`,
+jump 8 lessThan _main_i 10 	                 	
+end                        	                 	`,
 		},
 		{
 			name:  "All",
@@ -133,7 +136,8 @@ set _main_0 @return_0      	# 11	# Set variable to returned value
 print _main_0              	# 12	# Call to native function       	# println(foo(i))	
 print "\n"                 	# 13	# Call to native function       	# println(foo(i))	
 op add _main_i _main_i 1   	# 14	# Execute increment/decrement   	# i++            	
-jump 8 lessThan _main_i 10 	# 15	# Jump to start of loop         	                 	`,
+jump 8 lessThan _main_i 10 	# 15	# Jump to start of loop         	                 	
+end                        	# 16	# Trampoline back               	                 	`,
 		},
 	}
 	for _, test := range tests {

diff --git a/tests/stacked_function_test.go b/tests/stacked_function_test.go
@@ -271,6 +271,7 @@ op sub @stack @stack 4`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/stackless_function_test.go b/tests/stackless_function_test.go
@@ -232,6 +232,7 @@ jump 5 always`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/statement_test.go b/tests/statement_test.go
@@ -158,6 +158,15 @@ set _main_b distance
 set _main_c @this
 set _main_d core`,
 		},
+		{
+			name:  "MainReturn",
+			input: TestMain(`if x > 10 { return }; print(x)`),
+			output: `op greaterThan _main_0 _main_x 10
+jump 3 equal _main_0 1
+jump 4 always
+end
+print _main_x`,
+		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
@@ -170,6 +179,7 @@ set _main_d core`,
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/unit_control_test.go b/tests/unit_control_test.go
@@ -105,6 +105,7 @@ func TestUnitControl(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/tests/unit_test.go b/tests/unit_test.go
@@ -55,6 +55,7 @@ func TestUnit(t *testing.T) {
 				return
 			}
 
+			test.output = test.output + "\nend"
 			assert.Equal(t, test.output, strings.Trim(mlog, "\n"))
 		})
 	}

diff --git a/transpiler/main.go b/transpiler/main.go
@@ -196,6 +196,11 @@ func GolangToMLOG(input string, options Options) (string, error) {
 		return "", Err(ctx, "empty main function")
 	}
 
+	mainStatements = append(mainStatements, &MLOGTrampolineBack{
+		Stacked:  ctx.Value(contextOptions).(Options).Stacked,
+		Function: mainFuncName,
+	})
+
 	global.Functions = append(global.Functions, &Function{
 		Name:          mainFuncName,
 		Called:        true,

diff --git a/transpiler/type_native.go b/transpiler/type_native.go
@@ -34,7 +34,13 @@ type MLOGTrampolineBack struct {
 }
 
 func (m *MLOGTrampolineBack) PreProcess(ctx context.Context, global *Global, function *Function) error {
-	if m.Stacked != "" {
+	if m.Function == mainFuncName {
+		m.Statement = [][]Resolvable{
+			{
+				&Value{Value: "end"},
+			},
+		}
+	} else if m.Stacked != "" {
 		m.Statement = [][]Resolvable{
 			{
 				&Value{Value: "read"},