feat(formatter): Stricter token and scoping rules

3 files changed, 66 insertions, 17 deletions

diff --git a/README.md b/README.md
index a519417..961ab36 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,7 @@
 
 Let your code breathe!
 
-Iku is a grammar-based Go formatter that enforces consistent blank-line placement by AST node type.
+Iku is a grammar-based Go formatter that enforces consistent blank-line placement by statement and declaration type.
 
 ## Philosophy
 
@@ -13,10 +13,10 @@ Code structure should be visually apparent from its formatting. Iku groups state
 
 ## Rules
 
-1. **Same AST type means no blank line**: Consecutive statements of the same type stay together
-2. **Different AST type means blank line**: Transitions between statement types get visual separation
-3. **Scoped statements get blank lines**: `if`, `for`, `switch`, `select` always have blank lines before them
-4. **Top-level declarations are separated**: Functions, types, and variables at the package level get blank lines between them
+1. **Same type means no blank line**: Consecutive statements of the same type stay together
+2. **Different type means blank line**: Transitions between statement types get visual separation
+3. **Scoped constructs get blank lines**: `if`, `for`, `switch`, `select`, `func`, `type struct`, `type interface` always have blank lines around them
+4. **Declarations use token types**: `var`, `const`, `type`, `func`, `import` are distinguished by their keyword, not grouped as generic declarations
 
 ## How It Works
 
@@ -134,7 +134,10 @@ package main
 type Config struct {
     Name string
 }
+type ID int
+type Name string
 var defaultConfig = Config{}
+var x = 1
 func main() {
     run()
 }
@@ -149,7 +152,11 @@ type Config struct {
     Name string
 }
 
+type ID int
+type Name string
+
 var defaultConfig = Config{}
+var x = 1
 
 func main() {
     run()
@@ -160,6 +167,12 @@ func run() {
 }
 ```
 
+Notice how:
+- `type Config struct` is scoped (has braces), so it gets a blank line
+- `type ID int` and `type Name string` are unscoped type aliases, so they group together
+- `var defaultConfig` and `var x` are unscoped, so they group together
+- `func main()` and `func run()` are scoped, so each gets a blank line
+
 ### Switch Statements
 
 ```go
diff --git a/formatter.go b/formatter.go
index 8231bb7..2ac28b5 100644
--- a/formatter.go
+++ b/formatter.go
@@ -13,7 +13,7 @@ import (
 var (
 	closingBracePattern = regexp.MustCompile(`^\s*[\}\)]`)
 	openingBracePattern = regexp.MustCompile(`[\{\(]\s*$`)
-	caseLabelPattern    = regexp.MustCompile(`^\s*(case\s+.*|default\s*):\s*$`)
+	caseLabelPattern    = regexp.MustCompile(`^\s*(case\s|default\s*:)|(^\s+.*:\s*$)`)
 )
 
 func isCommentOnly(line string) bool {
@@ -128,6 +128,19 @@ func (f *Formatter) Format(source []byte) ([]byte, error) {
 	return f.rewrite(formatted, lineInfoMap), nil
 }
 
+func isGenDeclScoped(genDecl *ast.GenDecl) bool {
+	for _, spec := range genDecl.Specs {
+		if typeSpec, ok := spec.(*ast.TypeSpec); ok {
+			switch typeSpec.Type.(type) {
+			case *ast.StructType, *ast.InterfaceType:
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
 func (f *Formatter) buildLineInfo(fileSet *token.FileSet, file *ast.File) map[int]*lineInfo {
 	lineInfoMap := make(map[int]*lineInfo)
 	tokenFile := fileSet.File(file.Pos())
@@ -139,9 +152,22 @@ func (f *Formatter) buildLineInfo(fileSet *token.FileSet, file *ast.File) map[in
 	for _, declaration := range file.Decls {
 		startLine := tokenFile.Line(declaration.Pos())
 		endLine := tokenFile.Line(declaration.End())
-		typeName := reflect.TypeOf(declaration).String()
-		lineInfoMap[startLine] = &lineInfo{statementType: typeName, isTopLevel: true, isStartLine: true}
-		lineInfoMap[endLine] = &lineInfo{statementType: typeName, isTopLevel: true, isStartLine: false}
+		typeName := ""
+		isScoped := false
+
+		switch declarationType := declaration.(type) {
+		case *ast.GenDecl:
+			typeName = declarationType.Tok.String()
+			isScoped = isGenDeclScoped(declarationType)
+		case *ast.FuncDecl:
+			typeName = "func"
+			isScoped = true
+		default:
+			typeName = reflect.TypeOf(declaration).String()
+		}
+
+		lineInfoMap[startLine] = &lineInfo{statementType: typeName, isTopLevel: true, isScoped: isScoped, isStartLine: true}
+		lineInfoMap[endLine] = &lineInfo{statementType: typeName, isTopLevel: true, isScoped: isScoped, isStartLine: false}
 	}
 
 	ast.Inspect(file, func(node ast.Node) bool {
@@ -176,14 +202,22 @@ func (f *Formatter) processStatementList(tokenFile *token.File, statements []ast
 	for _, statement := range statements {
 		startLine := tokenFile.Line(statement.Pos())
 		endLine := tokenFile.Line(statement.End())
-		typeName := reflect.TypeOf(statement).String()
+		typeName := ""
 		isScoped := false
 
-		switch statement.(type) {
+		switch statementType := statement.(type) {
+		case *ast.DeclStmt:
+			if genericDeclaration, ok := statementType.Decl.(*ast.GenDecl); ok {
+				typeName = genericDeclaration.Tok.String()
+			} else {
+				typeName = reflect.TypeOf(statement).String()
+			}
 		case *ast.IfStmt, *ast.ForStmt, *ast.RangeStmt, *ast.SwitchStmt,
 			*ast.TypeSwitchStmt, *ast.SelectStmt, *ast.BlockStmt:
-
+			typeName = reflect.TypeOf(statement).String()
 			isScoped = true
+		default:
+			typeName = reflect.TypeOf(statement).String()
 		}
 
 		existingStart := lineInfoMap[startLine]
@@ -258,6 +292,7 @@ func (f *Formatter) rewrite(source []byte, lineInfoMap map[int]*lineInfo) []byte
 	previousType := ""
 	previousWasComment := false
 	previousWasTopLevel := false
+	previousWasScoped := false
 	insideRawString := false
 
 	for index, line := range lines {
@@ -302,12 +337,12 @@ func (f *Formatter) rewrite(source []byte, lineInfoMap map[int]*lineInfo) []byte
 		currentIsScoped := info != nil && info.isScoped
 
 		if len(result) > 0 && !previousWasOpenBrace && !isClosingBrace && !isCaseLabel {
-			if currentIsTopLevel && previousWasTopLevel {
+			if currentIsTopLevel && previousWasTopLevel && currentType != previousType {
 				if f.CommentMode == CommentsFollow && previousWasComment {
 				} else {
 					needsBlank = true
 				}
-			} else if currentIsScoped {
+			} else if info != nil && (currentIsScoped || previousWasScoped) {
 				if f.CommentMode == CommentsFollow && previousWasComment {
 				} else {
 					needsBlank = true
@@ -329,9 +364,9 @@ func (f *Formatter) rewrite(source []byte, lineInfoMap map[int]*lineInfo) []byte
 						nextIsTopLevel := nextInfo.isTopLevel
 						nextIsScoped := nextInfo.isScoped
 
-						if nextIsTopLevel && previousWasTopLevel {
+						if nextIsTopLevel && previousWasTopLevel && nextInfo.statementType != previousType {
 							needsBlank = true
-						} else if nextIsScoped {
+						} else if nextIsScoped || previousWasScoped {
 							needsBlank = true
 						} else if nextInfo.statementType != "" && previousType != "" && nextInfo.statementType != previousType {
 							needsBlank = true
@@ -352,9 +387,11 @@ func (f *Formatter) rewrite(source []byte, lineInfoMap map[int]*lineInfo) []byte
 		if info != nil {
 			previousType = info.statementType
 			previousWasTopLevel = info.isTopLevel
+			previousWasScoped = info.isScoped
 		} else if currentType != "" {
 			previousType = currentType
 			previousWasTopLevel = false
+			previousWasScoped = false
 		}
 	}
 
diff --git a/main.go b/main.go
index 2e6c7a0..e9a2014 100644
--- a/main.go
+++ b/main.go
@@ -13,7 +13,6 @@ import (
 )
 
 var version = "dev"
-
 var (
 	writeFlag    = flag.Bool("w", false, "write result to (source) file instead of stdout")
 	listFlag     = flag.Bool("l", false, "list files whose formatting differs from iku's")