Luke Wilson 2af696078d Improve Zig AST developer experience
- Convert string-based type discrimination to type-safe enums
  - ContainerKind (struct, enum, union, opaque)
  - LiteralKind (int, float, string, char)
  - LoopKind (for, while)
- Remove duplicate AST nodes (consolidated init lists, removed unused types)
- Add comprehensive helper functions for all AST constructions
- Implement formatters for all AST nodes (expressions, statements, types)
- Add typed literal constructors: IntLit, FloatLit, StringLit, CharLit
- Improve documentation and add deprecation notices

This makes the AST more intuitive and type-safe for developers.
2025-06-05 21:10:50 -05:00

603 lines
12 KiB
Go

package zig
import (
"fmt"
"io"
"strings"
)
type formatter struct {
w io.Writer
line int // 1-based
col int // 1-based, reset to 1 after newline
indent int // indentation level
}
// indentStr defines the string used for each indentation level (4 spaces).
const indentStr = " "
// writef writes formatted text to the underlying writer and updates line/col counters.
// It also handles indentation after newlines when appropriate.
func (f *formatter) writef(format string, a ...any) {
s := fmt.Sprintf(format, a...)
for i, r := range s {
if r == '\n' {
f.line++
f.col = 1
// After a newline, write indentation for the next line unless it's a closing brace or another newline.
if i+1 < len(s) && s[i+1] != '\n' && s[i+1] != '}' {
f.writeIndent()
}
} else {
if f.col == 0 {
f.col = 1
} else {
f.col++
}
}
}
if _, err := f.w.Write([]byte(s)); err != nil {
panic(err)
}
}
// writeIndent writes the current indentation level to the output.
// Call this at the start of a new line before writing statements or closing braces.
func (f *formatter) writeIndent() {
for i := 0; i < f.indent; i++ {
if _, err := f.w.Write([]byte(indentStr)); err != nil {
panic(err)
}
f.col += len(indentStr)
}
}
// Write is the entry point for formatting a Zig AST.
func Write(w io.Writer, root *Root) (err error) {
defer func() {
if r := recover(); r != nil {
if e, ok := r.(error); ok {
err = e
} else {
panic(r)
}
}
}()
sb := &strings.Builder{}
f := &formatter{w: sb, line: 1, col: 1, indent: 0}
if root.ContainerDocComment != "" {
f.writef("//! %s\n\n", root.ContainerDocComment)
}
for i, member := range root.ContainerMembers {
if member.Decl != nil {
// Only emit a leading newline before a function/global after the first declaration
if i > 0 {
f.writef("\n")
}
writeDecl(f, member.Decl)
}
}
out := sb.String()
if len(out) == 0 || out[len(out)-1] != '\n' {
out += "\n"
}
_, err = w.Write([]byte(out))
return err
}
// writeDecl emits a top-level declaration.
func writeDecl(f *formatter, decl Decl) {
switch d := decl.(type) {
case *FnDecl:
if d.Flags&FnExport != 0 {
f.writef("pub ")
}
f.writef("fn %s(", d.Name)
writeParams(f, d.Params)
f.writef(") ")
writeTypeExpr(f, d.ReturnType)
writeBlock(f, d.Body)
f.writef("\n")
case *GlobalVarDecl:
if d.Flags&GlobalVarConst != 0 {
f.writef("const %s = ", d.Name)
} else {
f.writef("var %s = ", d.Name)
}
writeExpr(f, d.Value)
f.writef(";\n")
case *ContainerDecl:
f.writef("struct ")
writeStructBody(f, d)
case *ErrorSetDecl:
writeExpr(f, d)
case *TestDecl:
f.writef("test ")
if d.Name != "" {
f.writef(`"%s"`, d.Name)
}
writeBlock(f, d.Block)
f.writef("\n")
case *UsingNamespaceDecl:
f.writef("usingnamespace ")
writeExpr(f, d.Expr)
f.writef(";\n")
case *ComptimeDecl:
f.writef("comptime")
writeBlock(f, d.Block)
f.writef("\n")
}
}
// writeParams emits function parameters, separated by commas.
func writeParams(f *formatter, params []*ParamDecl) {
for i, param := range params {
if i > 0 {
f.writef(", ")
}
if param.Name != "" {
f.writef("%s: ", param.Name)
}
writeTypeExpr(f, param.Type)
}
}
// writeTypeExpr emits a type expression.
func writeTypeExpr(f *formatter, typ TypeExpr) {
switch t := typ.(type) {
case *Identifier:
f.writef("%s", t.Name)
case *PrefixTypeExpr:
f.writef("%s", t.Op)
writeTypeExpr(f, t.Base)
case *ArrayTypeExpr:
f.writef("[")
writeExpr(f, t.Size)
f.writef("]")
writeTypeExpr(f, t.Elem)
case *SliceTypeExpr:
f.writef("[]")
writeTypeExpr(f, t.Elem)
case *ErrorUnionTypeExpr:
writeTypeExpr(f, t.ErrSet)
f.writef("!")
writeTypeExpr(f, t.Type)
case nil:
// nothing
default:
f.writef("%v", t)
}
}
// writeBlock emits a block, handling indentation for statements and the closing brace.
func writeBlock(f *formatter, block *Block) {
if block == nil {
f.writef(";")
return
}
f.writef(" {\n")
f.indent++
for _, stmt := range block.Stmts {
f.writeIndent()
writeStmt(f, stmt)
f.writef("\n")
}
f.indent--
f.writeIndent()
f.writef("}")
}
// writeStmt emits a statement. Indentation is handled by the caller (writeBlock).
func writeStmt(f *formatter, stmt Stmt) {
switch s := stmt.(type) {
case *ReturnStmt:
f.writef("return")
if s.Value != nil {
f.writef(" ")
writeExpr(f, s.Value)
}
f.writef(";")
case *ExprStmt:
writeExpr(f, s.Expr)
f.writef(";")
case *VarDeclStmt:
if s.Const {
f.writef("const ")
} else {
f.writef("var ")
}
for i, name := range s.Pattern.Names {
if i > 0 {
f.writef(", ")
}
f.writef("%s", name)
}
if s.Type != nil {
f.writef(": ")
writeTypeExpr(f, s.Type)
}
if s.Value != nil {
f.writef(" = ")
writeExpr(f, s.Value)
}
f.writef(";")
case *BlockStmt:
writeBlock(f, s.Block)
case *IfStmt:
f.writef("if (")
writeExpr(f, s.Cond)
f.writef(")")
// Handle payload if present
if s.Payload != nil {
f.writef(" ")
writePayload(f, s.Payload)
}
// Always write the then branch as a block
if block, ok := s.Then.(*BlockStmt); ok {
writeBlock(f, block.Block)
} else {
f.writef(" ")
writeStmt(f, s.Then)
}
if s.Else != nil {
f.writef(" else")
if block, ok := s.Else.(*BlockStmt); ok {
writeBlock(f, block.Block)
} else {
f.writef(" ")
writeStmt(f, s.Else)
}
}
case *LoopStmt:
if s.Kind == LoopWhile {
f.writef("while (")
if wp, ok := s.Prefix.(*WhilePrefix); ok {
writeExpr(f, wp.Cond)
if wp.Continue != nil {
f.writef(") : (")
writeExpr(f, wp.Continue)
}
}
f.writef(")")
// Always write the body as a block
if block, ok := s.Body.(*BlockStmt); ok {
writeBlock(f, block.Block)
} else {
f.writef(" ")
writeStmt(f, s.Body)
}
} else if s.Kind == LoopFor {
f.writef("for (")
if fp, ok := s.Prefix.(*ForPrefix); ok {
for i, arg := range fp.Args {
if i > 0 {
f.writef(", ")
}
writeExpr(f, arg.Expr)
if arg.From != nil {
if lit, ok := arg.From.(*Literal); ok && lit.Value == "" {
f.writef("..")
} else {
f.writef("...")
writeExpr(f, arg.From)
}
}
}
f.writef(")")
if fp.Payload != nil {
f.writef(" ")
writePayload(f, fp.Payload)
}
}
// Always write the body as a block
if block, ok := s.Body.(*BlockStmt); ok {
writeBlock(f, block.Block)
} else {
f.writef(" ")
writeStmt(f, s.Body)
}
}
case *DeferStmt:
if s.ErrDefer {
f.writef("errdefer")
} else {
f.writef("defer")
}
f.writef(" ")
writeStmt(f, s.Stmt)
case *BreakStmt:
f.writef("break")
if s.Label != "" {
f.writef(" :%s", s.Label)
}
if s.Value != nil {
f.writef(" ")
writeExpr(f, s.Value)
}
f.writef(";")
case *ContinueStmt:
f.writef("continue")
if s.Label != "" {
f.writef(" :%s", s.Label)
}
f.writef(";")
case *SwitchStmt:
f.writef("switch (")
writeExpr(f, s.Cond)
f.writef(") {\n")
f.indent++
for _, prong := range s.Prongs {
f.writeIndent()
writeSwitchProng(f, prong)
f.writef("\n")
}
f.indent--
f.writeIndent()
f.writef("}")
}
}
// writePayload emits a payload (|x|, |*x|, |*x, y|, etc).
func writePayload(f *formatter, payload *Payload) {
f.writef("|")
for i, name := range payload.Names {
if i > 0 {
f.writef(", ")
}
if payload.Pointers[i] {
f.writef("*")
}
f.writef("%s", name)
}
f.writef("|")
}
// writeSwitchProng emits a switch prong.
func writeSwitchProng(f *formatter, prong *SwitchProng) {
for i, c := range prong.Cases {
if i > 0 {
f.writef(", ")
}
if c.IsElse {
f.writef("else")
} else {
writeExpr(f, c.Expr)
if c.To != nil {
f.writef("...")
writeExpr(f, c.To)
}
}
}
f.writef(" => ")
// Check if the expression is actually a statement (like return or break)
if stmt, ok := prong.Expr.(Stmt); ok {
// If it's a block, write it directly without the leading space
if blockStmt, isBlock := stmt.(*BlockStmt); isBlock {
f.writef("{\n")
f.indent++
for _, s := range blockStmt.Block.Stmts {
f.writeIndent()
writeStmt(f, s)
f.writef("\n")
}
f.indent--
f.writeIndent()
f.writef("},")
} else {
// For single statements, write without the semicolon
switch s := stmt.(type) {
case *ReturnStmt:
f.writef("return")
if s.Value != nil {
f.writef(" ")
writeExpr(f, s.Value)
}
case *BreakStmt:
f.writef("break")
if s.Label != "" {
f.writef(" :%s", s.Label)
}
if s.Value != nil {
f.writef(" ")
writeExpr(f, s.Value)
}
case *ContinueStmt:
f.writef("continue")
if s.Label != "" {
f.writef(" :%s", s.Label)
}
case *ExprStmt:
writeExpr(f, s.Expr)
default:
writeStmt(f, stmt)
}
f.writef(",")
}
} else {
writeExpr(f, prong.Expr)
f.writef(",")
}
}
// writeExpr emits an expression.
func writeExpr(f *formatter, expr Expr) {
switch e := expr.(type) {
case *Identifier:
f.writef("%s", e.Name)
case *CallExpr:
writeExpr(f, e.Fun)
f.writef("(")
for i, arg := range e.Args {
if i > 0 {
f.writef(", ")
}
writeExpr(f, arg)
}
f.writef(")")
case *FieldAccessExpr:
writeExpr(f, e.Receiver)
f.writef(".%s", e.Field)
case *Literal:
switch e.Kind {
case LiteralString:
f.writef(`"%v"`, e.Value)
default:
f.writef("%v", e.Value)
}
case *InitListExpr:
if e.Empty {
f.writef(".{}")
} else if len(e.Values) > 0 {
if len(e.Values) == 1 {
f.writef(".{")
writeExpr(f, e.Values[0])
f.writef("}")
} else {
f.writef(".{ ")
for i, v := range e.Values {
if i > 0 {
f.writef(", ")
}
writeExpr(f, v)
}
f.writef(" }")
}
}
case *ContainerDecl:
switch e.Kind {
case ContainerStruct:
f.writef("struct ")
writeStructBody(f, e)
case ContainerEnum:
f.writef("enum ")
writeStructBody(f, e)
case ContainerUnion:
f.writef("union ")
writeStructBody(f, e)
case ContainerOpaque:
f.writef("opaque ")
writeStructBody(f, e)
default:
panic("unknown container kind")
}
case *TryExpr:
f.writef("try ")
writeExpr(f, e.Expr)
case *BinaryExpr:
writeExpr(f, e.Left)
f.writef(" %s ", e.Op)
writeExpr(f, e.Right)
case *UnaryExpr:
f.writef("%s", e.Op)
writeExpr(f, e.Expr)
case *IndexExpr:
writeExpr(f, e.Receiver)
f.writef("[")
writeExpr(f, e.Index)
f.writef("]")
case *UnreachableExpr:
f.writef("unreachable")
case *ErrorSetDecl:
f.writef("error{\n")
f.indent++
for _, name := range e.Names {
f.writeIndent()
f.writef("%s,\n", name)
}
f.indent--
f.writeIndent()
f.writef("}")
case *GroupedExpr:
f.writef("(")
writeExpr(f, e.Expr)
f.writef(")")
case *ComptimeExpr:
f.writef("comptime ")
writeExpr(f, e.Expr)
case *NosuspendExpr:
f.writef("nosuspend ")
writeExpr(f, e.Expr)
case *AsyncExpr:
f.writef("async ")
writeExpr(f, e.Expr)
case *AwaitExpr:
f.writef("await ")
writeExpr(f, e.Expr)
case *ResumeExpr:
f.writef("resume ")
writeExpr(f, e.Expr)
case *DotAsteriskExpr:
writeExpr(f, e.Receiver)
f.writef(".*")
case *DotQuestionExpr:
writeExpr(f, e.Receiver)
f.writef(".?")
case *IfExpr:
f.writef("if (")
writeExpr(f, e.Cond)
f.writef(")")
if e.Payload != nil {
f.writef(" ")
writePayload(f, e.Payload)
}
f.writef(" ")
writeExpr(f, e.Then)
if e.Else != nil {
f.writef(" else ")
writeExpr(f, e.Else)
}
case *SwitchExpr:
f.writef("switch (")
writeExpr(f, e.Cond)
f.writef(") {\n")
f.indent++
for _, prong := range e.Prongs {
f.writeIndent()
writeSwitchProng(f, prong)
f.writef("\n")
}
f.indent--
f.writeIndent()
f.writef("}")
}
}
// writeStructBody emits the body of a struct/union/enum/opaque declaration.
func writeStructBody(f *formatter, decl *ContainerDecl) {
f.writef("{\n")
f.indent++
for _, member := range decl.Fields {
if member.Field != nil {
// Field or const
if member.Field.Type == nil && member.Field.Value != nil {
// const field
f.writeIndent()
f.writef("const %s = ", member.Field.Name)
writeExpr(f, member.Field.Value)
f.writef(";\n")
} else {
// regular field
f.writeIndent()
f.writef("%s: ", member.Field.Name)
writeTypeExpr(f, member.Field.Type)
if member.Field.Value != nil {
f.writef(" = ")
writeExpr(f, member.Field.Value)
}
f.writef(",\n")
}
} else if member.Decl != nil {
// Method or nested decl
f.writef("\n")
f.writeIndent()
writeDecl(f, member.Decl)
}
}
f.indent--
f.writeIndent()
f.writef("}")
}