Nearly complete Zig AST

internal/zig/zig.go

@@ -6,32 +6,30 @@ import (
 )
 
 type formatter struct {
-	w io.Writer
+	io.Writer
 }
 
 func (f *formatter) WriteString(s string) (n int, err error) {
-	return f.w.Write([]byte(s))
+	return f.Write([]byte(s))
 }
 
 func (f *formatter) Writef(format string, a ...any) (err error) {
-	_, err = f.w.Write(fmt.Appendf(nil, format, a...))
+	_, err = fmt.Fprintf(f, format, a...)
 	return err
 }
 
 func Write(w io.Writer, root *Root) error {
-	f := &formatter{
-		w: w,
-	}
+	f := &formatter{Writer: w}
 
 	if root.ContainerDocComment != "" {
-		err := f.Writef("//! %s\n\n", root.ContainerDocComment)
-		if err != nil {
+		if err := f.Writef("//! %s\n\n", root.ContainerDocComment); err != nil {
 			return err
 		}
 	}
 	for _, member := range root.ContainerMembers {
-		for _, decl := range member.Decls {
-			if err := writeDecl(f, decl); err != nil {
+		// Only handle Decl for now (fields not needed for hello world)
+		if member.Decl != nil {
+			if err := writeDecl(f, member.Decl); err != nil {
 				return err
 			}
 		}
@@ -39,68 +37,90 @@ func Write(w io.Writer, root *Root) error {
 	return nil
 }
 
-func writeDecl(f *formatter, decl Decl) (err error) {
-	switch typ := decl.(type) {
+func writeDecl(f *formatter, decl Decl) error {
+	switch fn := decl.(type) {
 	case *FnDecl:
-		if err = f.Writef("fn %s(", typ.Name); err != nil {
+		if err := f.Writef("fn %s(", fn.Name); err != nil {
 			return err
 		}
-		if err = writeParams(f, typ.Params); err != nil {
+		if err := writeParams(f, fn.Params); err != nil {
 			return err
 		}
-		if err = f.Writef(") %s", typ.ReturnType); err != nil {
+		if _, err := f.WriteString(") "); err != nil {
 			return err
 		}
-		if err = writeBlock(f, typ.Body); err != nil {
+		if err := writeTypeExpr(f, fn.ReturnType); err != nil {
+			return err
+		}
+		if err := writeBlock(f, fn.Body); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func writeParams(f *formatter, params []*ParamDecl) (err error) {
-	for _, param := range params {
-		if param.Name != "" {
-			if err = f.Writef("%s: ", param.Name); err != nil {
+func writeParams(f *formatter, params []*ParamDecl) error {
+	for i, param := range params {
+		if i > 0 {
+			if _, err := f.WriteString(", "); err != nil {
 				return err
 			}
 		}
-		if err = f.Writef("%s", param.Type); err != nil {
+		if param.Name != "" {
+			if _, err := f.WriteString(param.Name); err != nil {
+				return err
+			}
+			if _, err := f.WriteString(": "); err != nil {
+				return err
+			}
+		}
+		if err := writeTypeExpr(f, param.Type); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func writeBlock(f *formatter, block *Block) (err error) {
-	if block == nil {
-		if _, err = f.WriteString(";"); err != nil {
-			return err
-		}
+func writeTypeExpr(f *formatter, typ TypeExpr) error {
+	switch t := typ.(type) {
+	case *Identifier:
+		_, err := f.WriteString(t.Name)
+		return err
+	case nil:
 		return nil
+	default:
+		// fallback: print as string
+		_, err := f.WriteString(fmt.Sprintf("%v", t))
+		return err
 	}
-	if err = f.Writef(" {\n"); err != nil {
+}
+
+func writeBlock(f *formatter, block *Block) error {
+	if block == nil {
+		_, err := f.WriteString(";")
+		return err
+	}
+	if _, err := f.WriteString(" {\n"); err != nil {
 		return err
 	}
 	for _, stmt := range block.Stmts {
-		if err = writeStmt(f, stmt); err != nil {
+		if err := writeStmt(f, stmt); err != nil {
 			return err
 		}
-		// Should this be the job of the formatter?
-		if _, err = f.WriteString("\n"); err != nil {
+		if _, err := f.WriteString("\n"); err != nil {
 			return err
 		}
 	}
-	if err = f.Writef("}\n"); err != nil {
+	if _, err := f.WriteString("}\n"); err != nil {
 		return err
 	}
 	return nil
 }
 
-func writeStmt(f *formatter, stmt Stmt) (err error) {
+func writeStmt(f *formatter, stmt Stmt) error {
 	switch stmt.(type) {
 	case *ReturnStmt:
-		if _, err = f.WriteString("return;"); err != nil {
+		if _, err := f.WriteString("return;"); err != nil {
 			return err
 		}
 	}