From b62bf59c21fb02e93e8c080df672da7a916d5410 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20D=C3=A4schle?= <daniel.daeschle@gmail.com>
Date: Wed, 8 Jul 2020 15:17:28 +0200
Subject: [PATCH] all: if-is smartcast part 1 (#5730)

---
 vlib/v/ast/ast.v         |  5 ++++-
 vlib/v/ast/str.v         |  2 +-
 vlib/v/checker/checker.v | 23 +++++++++++++++++++++++
 vlib/v/gen/cgen.v        | 14 ++++++++++++++
 vlib/v/gen/js/js.v       | 22 +++++++++++-----------
 vlib/v/parser/if.v       |  6 +++++-
 6 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 2a167f8b9a..7d70520f5a 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -425,10 +425,10 @@ pub mut:
 pub struct IfExpr {
 pub:
 	tok_kind token.Kind
-	branches []IfBranch
 	left     Expr // `a` in `a := if ...`
 	pos      token.Position
 pub mut:
+	branches []IfBranch
 	is_expr  bool
 	typ      table.Type
 	has_else bool
@@ -439,7 +439,10 @@ pub:
 	cond     Expr
 	stmts    []Stmt
 	pos      token.Position
+	body_pos token.Position
 	comments []Comment
+pub mut:
+	smartcast bool // should only be true if cond is `x is sumtype`, it will be set in checker - if_expr
 }
 
 pub struct LockExpr {
diff --git a/vlib/v/ast/str.v b/vlib/v/ast/str.v
index acf2575e30..b9f2dc7e61 100644
--- a/vlib/v/ast/str.v
+++ b/vlib/v/ast/str.v
@@ -286,7 +286,7 @@ pub fn (node Stmt) str() string {
 					}
 				}
 				out += left.str()
-				if i < it.left.len - 1 {
+				if i < node.left.len - 1 {
 					out += ','
 				}
 			}
diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 445bd14ef3..3b50292427 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -2623,6 +2623,29 @@ pub fn (mut c Checker) if_expr(mut node ast.IfExpr) table.Type {
 				c.error('non-bool type `$typ_sym.name` used as if condition', branch.pos)
 			}
 		}
+		// smartcast sumtypes when using `is`
+		if branch.cond is ast.InfixExpr {
+			infix := branch.cond as ast.InfixExpr
+			if infix.op == .key_is && infix.left is ast.Ident && infix.right is ast.Type {
+				left_expr := infix.left as ast.Ident
+				right_expr := infix.right as ast.Type
+				if left_expr.kind == .variable {
+					// Register shadow variable or `as` variable with actual type
+					left_sym := c.table.get_type_symbol(infix.left_type)
+					if left_sym.kind == .sum_type {
+						mut scope := c.file.scope.innermost(branch.body_pos.pos)
+						scope.register('it', ast.Var{
+							name: 'it'
+							typ: right_expr.typ.to_ptr()
+							pos: left_expr.pos
+							is_used: true
+							is_mut: left_expr.is_mut
+						})
+						node.branches[i].smartcast = true
+					}
+				}
+			}
+		}
 		c.stmts(branch.stmts)
 		if expr_required {
 			if branch.stmts.len > 0 && branch.stmts[branch.stmts.len - 1] is ast.ExprStmt {
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index bc612793b2..d9893d201f 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -2311,6 +2311,20 @@ fn (mut g Gen) if_expr(node ast.IfExpr) {
 		} else if i == node.branches.len - 1 && node.has_else {
 			g.writeln('} else {')
 		}
+		if branch.smartcast && branch.stmts.len > 0 {
+			infix := branch.cond as ast.InfixExpr
+			right_type := infix.right as ast.Type
+			left_type := infix.left_type
+			it_type := g.typ(right_type.typ)
+			g.write('\t$it_type* it = ($it_type*)')
+			g.expr(infix.left)
+			if left_type.is_ptr() {
+				g.write('->')
+			} else {
+				g.write('.')
+			}
+			g.writeln('obj;')
+		}
 		g.stmts(branch.stmts)
 	}
 	if is_guard {
diff --git a/vlib/v/gen/js/js.v b/vlib/v/gen/js/js.v
index d5e41f0268..8eb01b1849 100644
--- a/vlib/v/gen/js/js.v
+++ b/vlib/v/gen/js/js.v
@@ -663,27 +663,27 @@ fn (mut g JsGen) gen_assert_stmt(a ast.AssertStmt) {
 	g.writeln('}')
 }
 
-fn (mut g JsGen) gen_assign_stmt(it ast.AssignStmt) {
-	if it.left.len > it.right.len {
+fn (mut g JsGen) gen_assign_stmt(stmt ast.AssignStmt) {
+	if stmt.left.len > stmt.right.len {
 		// multi return
 		g.write('const [')
-		for i, left in it.left {
+		for i, left in stmt.left {
 			if !left.is_blank_ident() {
 				g.expr(left)
 			}
-			if i < it.left.len - 1 {
+			if i < stmt.left.len - 1 {
 				g.write(', ')
 			}
 		}
 		g.write('] = ')
-		g.expr(it.right[0])
+		g.expr(stmt.right[0])
 		g.writeln(';')
 	} else {
 		// `a := 1` | `a,b := 1,2`
-		for i, left in it.left {
-			mut op := it.op
-			if it.op == .decl_assign { op = .assign }
-			val := it.right[i]
+		for i, left in stmt.left {
+			mut op := stmt.op
+			if stmt.op == .decl_assign { op = .assign }
+			val := stmt.right[i]
 			mut is_mut := false
 			if left is ast.Ident {
 				ident := left as ast.Ident
@@ -698,13 +698,13 @@ fn (mut g JsGen) gen_assign_stmt(it ast.AssignStmt) {
 				}
 			}
 
-			mut styp := g.typ(it.left_types[i])
+			mut styp := g.typ(stmt.left_types[i])
 
 			if !g.inside_loop && styp.len > 0 {
 				g.doc.gen_typ(styp)
 			}
 
-			if it.op == .decl_assign {
+			if stmt.op == .decl_assign {
 				if g.inside_loop || is_mut {
 					g.write('let ')
 				} else {
diff --git a/vlib/v/parser/if.v b/vlib/v/parser/if.v
index 54a01e7781..2a1fe29f2e 100644
--- a/vlib/v/parser/if.v
+++ b/vlib/v/parser/if.v
@@ -31,9 +31,11 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 				has_else = true
 				p.inside_if = false
 				end_pos := p.prev_tok.position()
+				body_pos := p.tok.position()
 				branches << ast.IfBranch{
 					stmts: p.parse_block()
 					pos: start_pos.extend(end_pos)
+					body_pos: body_pos.extend(p.tok.position())
 					comments: comments
 				}
 				comments = []
@@ -63,7 +65,8 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 			cond = p.expr(0)
 		}
 		end_pos := p.prev_tok.position()
-		p.inside_if = false
+		body_pos := p.tok.position()
+		p.inside_if = false		
 		stmts := p.parse_block()
 		if is_or {
 			p.close_scope()
@@ -72,6 +75,7 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 			cond: cond
 			stmts: stmts
 			pos: start_pos.extend(end_pos)
+			body_pos: body_pos.extend(p.tok.position())
 			comments: comments
 		}
 		comments = []