Merge pull request #214 from x-motemen/struct-method

Support struct method declaration

session.go

@@ -224,25 +224,28 @@ func (s *Session) evalStmt(in string) error {
 	}
 
 	enclosingFunc := f.Scope.Lookup("F").Decl.(*ast.FuncDecl)
-	stmts := enclosingFunc.Body.List
-
-	if len(stmts) > 0 {
-		debugf("evalStmt :: %s", showNode(s.fset, stmts))
-		lastStmt := stmts[len(stmts)-1]
-		if assign, ok := lastStmt.(*ast.AssignStmt); ok {
-			if s := buildPrintStmt(assign.Lhs); s != nil {
-				stmts = append(stmts, s)
+
+	debugf("evalStmt :: %s", showNode(s.fset, enclosingFunc.Body.List))
+	var stmts []ast.Stmt
+
+	for _, stmt := range enclosingFunc.Body.List {
+		switch stmt := stmt.(type) {
+		case *ast.AssignStmt:
+			if stmt := buildPrintStmt(stmt.Lhs); stmt != nil {
+				stmts = append(stmts, stmt)
 			}
-		}
-		if decl, ok := lastStmt.(*ast.DeclStmt); ok {
-			if decl, ok := decl.Decl.(*ast.GenDecl); ok {
-				if s := buildPrintStmtOfDecl(decl); s != nil {
-					stmts = append(stmts, s)
+		case *ast.DeclStmt:
+			if decl, ok := stmt.Decl.(*ast.GenDecl); ok {
+				if decl.Tok == token.TYPE {
+					s.file.Decls = append(s.file.Decls, decl)
+					continue
+				} else if stmt := buildPrintStmtOfDecl(decl); stmt != nil {
+					stmts = append(stmts, stmt)
 				}
 			}
 		}
+		s.appendStatements(stmt)
 	}
-
 	s.appendStatements(stmts...)
 
 	return nil

session_test.go

@@ -303,6 +303,39 @@ func TestSessionEval_MultipleValues(t *testing.T) {
 	assert.Equal(t, "", stderr.String())
 }
 
+func TestSessionEval_Struct(t *testing.T) {
+	stdout, stderr := new(bytes.Buffer), new(bytes.Buffer)
+	s, err := NewSession(stdout, stderr)
+	defer s.Clear()
+	require.NoError(t, err)
+
+	codes := []string{
+		`type X struct { v int }`,
+		`func (x *X) add(v int) { x.v += v }`,
+		`var x X`,
+		`x`,
+		`x.add(1)`,
+		`x`,
+		`x.add(2)`,
+		`x`,
+		`type Y X; type Z Y;`,
+		`func (z *Z) sub(v int) { z.v -= v }`,
+		`var z Z`,
+		`z.sub(3)`,
+		`z`,
+	}
+
+	for _, code := range codes {
+		_ = s.Eval(code)
+	}
+
+	assert.Contains(t, stdout.String(), `main.X{v:0}
+main.X{v:1}
+main.X{v:3}
+main.Z{v:-3}`)
+	assert.Equal(t, ``, stderr.String())
+}
+
 func TestSessionEval_Func(t *testing.T) {
 	stdout, stderr := new(bytes.Buffer), new(bytes.Buffer)
 	s, err := NewSession(stdout, stderr)