Hoist expression out of switch tag

dispatchrun · Sep 20, 2023 · b7c305c · b7c305c
1 parent 3aee2ef
commit b7c305c
Show file tree

Hide file tree

Showing 3 changed files with 405 additions and 287 deletions.
diff --git a/compiler/desugar.go b/compiler/desugar.go
@@ -362,22 +362,34 @@ func (d *desugarer) desugar(stmt ast.Stmt, breakTo, continueTo, userLabel *ast.I
 
 	case *ast.SwitchStmt:
 		// Rewrite switch statements:
-		// - `switch init; tag { case A: ... case B: ... }` => `{ init; if _tag := tag; _tag == A { ... } else if _tag == B { ... }`
-		// - `switch { case A: ... case B: ... default: ... } => `if A { ... } else if B { ... } else { ... }`
+		// - `switch init; tag { ... }` => `{ init; _tag := tag; switch _tag { ... }`
 		switchLabel := d.newLabel()
 		if userLabel != nil {
 			d.addUserLabel(userLabel, switchLabel)
 		}
+		var prologue []ast.Stmt
+		if s.Init != nil {
+			prologue = []ast.Stmt{s.Init}
+		}
+		if s.Tag != nil {
+			tmp := d.newVar(d.info.TypeOf(s.Tag))
+			prologue = append(prologue, &ast.AssignStmt{
+				Lhs: []ast.Expr{tmp},
+				Tok: token.DEFINE,
+				Rhs: []ast.Expr{s.Tag},
+			})
+			s.Tag = tmp
+		}
 		// TODO: hoist each CaseClause.Cond out from SwitchStmt.Body so expressions can be desugared
+		prologue = d.desugarList(prologue, nil, nil)
 		stmt = &ast.LabeledStmt{
 			Label: switchLabel,
 			Stmt: &ast.SwitchStmt{
 				Tag:  s.Tag,
 				Body: d.desugar(s.Body, switchLabel, continueTo, nil).(*ast.BlockStmt),
 			},
 		}
-		if s.Init != nil {
-			prologue := d.desugarList([]ast.Stmt{s.Init}, nil, nil)
+		if len(prologue) > 0 {
 			stmt = &ast.BlockStmt{List: append(prologue, stmt)}
 		}
 

diff --git a/compiler/desugar_test.go b/compiler/desugar_test.go
@@ -609,13 +609,16 @@ default:
 	default:
 		_v0 = 3
 	}
-	switch _v0 {
-	case 1:
-		foo
-	case 2:
-		bar
-	case 3:
-		baz
+	{
+		_v1 := _v0
+		switch _v1 {
+		case 1:
+			foo
+		case 2:
+			bar
+		case 3:
+			baz
+		}
 	}
 }
 `,
@@ -640,12 +643,15 @@ label:
 	case <-b:
 		_v0 = 2
 	}
-_l0:
-	switch _v0 {
-	case 1:
-		break _l0
-	case 2:
-		break _l0
+	{
+		_v1 := _v0
+	_l0:
+		switch _v1 {
+		case 1:
+			break _l0
+		case 2:
+			break _l0
+		}
 	}
 }
 `,
@@ -663,7 +669,8 @@ default:
 			expect: `
 {
 	foo := bar
-	switch foo {
+	_v0 := foo
+	switch _v0 {
 	case 1:
 		bar
 	default:
@@ -683,11 +690,14 @@ default:
 }
 `,
 			expect: `
-switch foo {
-case 1:
-	bar
-default:
-	baz
+{
+	_v0 := foo
+	switch _v0 {
+	case 1:
+		bar
+	default:
+		baz
+	}
 }
 `,
 		},
@@ -800,7 +810,6 @@ l1:
 					continue l1
 				}
 		}
-	
 	}
 `,
 			defs: map[string]types.Object{
@@ -823,40 +832,43 @@ _l0:
 			default:
 				_v0 = 2
 			}
-		_l1:
-			switch _v0 {
-			case 1:
-				break _l1
-				break _l0
-				continue _l0
-			case 2:
-				{
-					_v1 := a
-				_l2:
-					switch _v1.(type) {
-					case int:
-						break _l2
-						break _l1
-						break _l0
-						continue _l0
-					_l3:
-						switch {
-						default:
-							break _l3
+			{
+				_v1 := _v0
+			_l1:
+				switch _v1 {
+				case 1:
+					break _l1
+					break _l0
+					continue _l0
+				case 2:
+					{
+						_v2 := a
+					_l2:
+						switch _v2.(type) {
+						case int:
 							break _l2
 							break _l1
 							break _l0
 							continue _l0
+						_l3:
+							switch {
+							default:
+								break _l3
+								break _l2
+								break _l1
+								break _l0
+								continue _l0
+							}
 						}
 					}
-				}
-			_l4:
-				for {
-					break _l4
-					break _l1
-					break _l0
-					continue _l4
-					continue _l0
+				_l4:
+					for {
+						break _l4
+						break _l1
+						break _l0
+						continue _l4
+						continue _l0
+					}
 				}
 			}
 		}