Added 2 tests and asthelpers.go

internal/zig/zig.go

@@ -3,6 +3,7 @@ package zig
 import (
 	"fmt"
 	"io"
+	"strings"
 )
 
 type formatter struct {
@@ -15,9 +16,9 @@ type formatter struct {
 // 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.
+// 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) {
+func (f *formatter) writef(format string, a ...any) {
 	s := fmt.Sprintf(format, a...)
 	for i, r := range s {
 		if r == '\n' {
@@ -62,37 +63,52 @@ func Write(w io.Writer, root *Root) (err error) {
 			}
 		}
 	}()
-	f := &formatter{w: w, line: 1, col: 1, indent: 0}
+	sb := &strings.Builder{}
+	f := &formatter{w: sb, line: 1, col: 1, indent: 0}
 
 	if root.ContainerDocComment != "" {
-		f.Writef("//! %s\n\n", root.ContainerDocComment)
+		f.writef("//! %s\n\n", root.ContainerDocComment)
 	}
-	for _, member := range root.ContainerMembers {
-		// Only handle Decl for now (fields not needed for hello world)
+	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)
 		}
 	}
-	return nil
+	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:
-		f.Writef("\nfn %s(", d.Name)
+		if d.Flags&FnExport != 0 {
+			f.writef("pub ")
+		}
+		f.writef("fn %s(", d.Name)
 		writeParams(f, d.Params)
-		f.Writef(") ")
+		f.writef(") ")
 		writeTypeExpr(f, d.ReturnType)
 		writeBlock(f, d.Body)
 	case *GlobalVarDecl:
-		if d.Const {
-			f.Writef("const %s = ", d.Name)
+		if d.Flags&GlobalVarConst != 0 {
+			f.writef("const %s = ", d.Name)
 		} else {
-			f.Writef("var %s = ", d.Name)
+			f.writef("var %s = ", d.Name)
 		}
 		writeExpr(f, d.Value)
-		f.Writef(";\n")
+		f.writef(";\n")
+	case *ContainerDecl:
+		f.writef("struct ")
+		writeStructBody(f, d)
 	}
 }
 
@@ -100,10 +116,10 @@ func writeDecl(f *formatter, decl Decl) {
 func writeParams(f *formatter, params []*ParamDecl) {
 	for i, param := range params {
 		if i > 0 {
-			f.Writef(", ")
+			f.writef(", ")
 		}
 		if param.Name != "" {
-			f.Writef("%s: ", param.Name)
+			f.writef("%s: ", param.Name)
 		}
 		writeTypeExpr(f, param.Type)
 	}
@@ -113,48 +129,111 @@ func writeParams(f *formatter, params []*ParamDecl) {
 func writeTypeExpr(f *formatter, typ TypeExpr) {
 	switch t := typ.(type) {
 	case *Identifier:
-		f.Writef("%s", t.Name)
+		f.writef("%s", t.Name)
+	case *PrefixTypeExpr:
+		f.writef("%s", t.Op)
+		writeTypeExpr(f, t.Base)
 	case nil:
 		// nothing
 	default:
-		f.Writef("%v", t)
+		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(";")
+		f.writef(";")
 		return
 	}
-	f.Writef(" {\n")
-	f.indent++ // Increase indentation for block contents.
-	for i, stmt := range block.Stmts {
-		f.writeIndent() // Indent each statement.
+	f.writef(" {\n")
+	f.indent++
+	for _, stmt := range block.Stmts {
+		f.writeIndent()
 		writeStmt(f, stmt)
-		if i < len(block.Stmts)-1 {
-			f.Writef("\n")
-		}
+		f.writef("\n")
 	}
-	f.indent-- // Decrease indentation before closing brace.
-	f.Writef("\n")
-	f.writeIndent() // Indent the closing brace.
-	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")
+		f.writef("return")
 		if s.Value != nil {
-			f.Writef(" ")
+			f.writef(" ")
 			writeExpr(f, s.Value)
 		}
-		f.Writef(";")
+		f.writef(";")
 	case *ExprStmt:
 		writeExpr(f, s.Expr)
-		f.Writef(";")
+		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)
+			}
+		}
 	}
 }
 
@@ -162,34 +241,90 @@ func writeStmt(f *formatter, stmt Stmt) {
 func writeExpr(f *formatter, expr Expr) {
 	switch e := expr.(type) {
 	case *Identifier:
-		f.Writef("%s", e.Name)
+		f.writef("%s", e.Name)
 	case *CallExpr:
 		writeExpr(f, e.Fun)
-		f.Writef("(")
+		f.writef("(")
 		for i, arg := range e.Args {
 			if i > 0 {
-				f.Writef(", ")
+				f.writef(", ")
 			}
 			writeExpr(f, arg)
 		}
-		f.Writef(")")
+		f.writef(")")
 	case *FieldAccessExpr:
 		writeExpr(f, e.Receiver)
-		f.Writef(".%s", e.Field)
+		f.writef(".%s", e.Field)
 	case *Literal:
 		switch e.Kind {
 		case "string":
-			f.Writef("%q", e.Value)
+			f.writef(`"%v"`, e.Value)
 		default:
-			f.Writef("%v", e.Value)
+			f.writef("%v", e.Value)
 		}
 	case *InitListExpr:
 		if e.Empty {
-			f.Writef(".{}")
-		} else {
-			f.Writef(".{")
-			// TODO
-			f.Writef("}")
+			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("}")
+}