331 lines
6.9 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)
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)
}
}
// 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 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(")")
// 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 == "while" {
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)
}
}
}
}
// 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 "string":
f.writef(`"%v"`, e.Value)
default:
f.writef("%v", e.Value)
}
case *InitListExpr:
if e.Empty {
f.writef(".{}")
} else if len(e.Values) > 0 {
f.writef(".{")
for i, v := range e.Values {
if i > 0 {
f.writef(", ")
}
writeExpr(f, v)
}
f.writef("}")
}
case *ContainerDecl:
if e.Kind == "struct" {
f.writef("struct ")
writeStructBody(f, e)
} else {
panic("not implemented: " + e.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)
}
}
// 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.writef("\n")
}
}
f.indent--
f.writeIndent()
f.writef("}")
}