ref counted strings

compiler/Generator.cs

@@ -5,36 +5,35 @@ namespace Compiler;
 
 public class Generator
 {
-    public static string Emit(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph, bool compileLib)
+    public static string Emit(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph, string? entryPoint)
     {
-        return new Generator(functions, moduleGraph, compileLib).Emit();
+        return new Generator(functions, moduleGraph, entryPoint).Emit();
     }
 
-    private Generator(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph, bool compileLib)
+    private Generator(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph, string? entryPoint)
     {
         this.functions = functions;
         this.moduleGraph = moduleGraph;
-        this.compileLib = compileLib;
+        this.entryPoint = entryPoint;
     }
 
     private readonly List<TypedNodeDefinitionFunc> functions;
     private readonly ModuleGraph moduleGraph;
-    private readonly bool compileLib;
+    private readonly string? entryPoint;
     private IndentedTextWriter writer = new();
     private HashSet<NubType> referencedTypes = new();
     private readonly HashSet<NubType> emittedTypes = new();
+    private readonly Stack<Scope> scopes = new();
+    private int tmpNameIndex = 0;
 
     private string Emit()
     {
-        if (!compileLib)
+        if (entryPoint != null)
         {
-            if (!moduleGraph.TryResolveIdentifier("main", "main", true, out var info))
-                throw new UnreachableException("func main::main() is not defined. This should have been caught earlier");
-
             writer.WriteLine($$"""
                              int main(int argc, char *argv[])
                              {
-                                return {{info.MangledName}}();
+                                return {{entryPoint}}();
                              }
                              """);
 
@@ -57,7 +56,9 @@ public class Generator
             writer.WriteLine("{");
             using (writer.Indent())
             {
+                PushScope();
                 EmitStatement(function.Body);
+                PopScope();
             }
             writer.WriteLine("}");
             writer.WriteLine();
@@ -67,12 +68,36 @@ public class Generator
 
         writer = new IndentedTextWriter();
 
+        foreach (var module in moduleGraph.GetModules())
+        {
+            foreach (var (name, info) in module.GetIdentifiers())
+            {
+                if (info.Source == Module.DefinitionSource.Internal || info.Exported)
+                {
+                    if (info.Source == Module.DefinitionSource.Imported || info.Extern)
+                        writer.Write("extern ");
+                    else if (info.Source == Module.DefinitionSource.Internal && !info.Extern && !info.Exported)
+                        writer.Write("static ");
+
+                    if (info.Type is NubTypeFunc fn)
+                        writer.WriteLine($"{CType(fn.ReturnType, info.MangledName)}({string.Join(", ", fn.Parameters.Select(p => CType(p)))});");
+                    else
+                        writer.WriteLine($"{CType(info.Type, info.MangledName)};");
+                }
+            }
+        }
+
+        var declarations = writer.ToString();
+
+        writer = new IndentedTextWriter();
+
         writer.WriteLine("""
                          #include <float.h>
                          #include <stdarg.h>
                          #include <stddef.h>
                          #include <stdbool.h>
                          #include <stdint.h>
+                         #include <stdlib.h>
 
                          """);
 
@@ -83,25 +108,9 @@ public class Generator
             referencedTypes.Remove(type);
         }
 
-        foreach (var module in moduleGraph.GetModules())
-        {
-            foreach (var (name, info) in module.GetIdentifiers())
-            {
-                if (info.Source == Module.DefinitionSource.Imported || info.Extern)
-                    writer.Write("extern ");
-                else if (info.Source == Module.DefinitionSource.Internal && !info.Extern && !info.Exported)
-                    writer.Write("static ");
-
-                if (info.Type is NubTypeFunc fn)
-                    writer.WriteLine($"{CType(fn.ReturnType, info.MangledName)}({string.Join(", ", fn.Parameters.Select(p => CType(p)))});");
-                else
-                    writer.WriteLine($"{CType(info.Type, info.MangledName)};");
-            }
-        }
-
         var header = writer.ToString();
 
-        return $"{header}\n{implementations}";
+        return $"{header}\n{declarations}\n{implementations}";
     }
 
     private void EmitTypeDefinitionIfNotEmitted(NubType type)
@@ -121,6 +130,7 @@ public class Generator
                 {
                     writer.WriteLine("char *data;");
                     writer.WriteLine("size_t length;");
+                    writer.WriteLine("uint32_t ref;");
                 }
                 writer.WriteLine("};");
                 writer.WriteLine();
@@ -211,6 +221,9 @@ public class Generator
 
     private void EmitStatement(TypedNodeStatement node)
     {
+        if (scopes.Peek().Unreachable)
+            return;
+
         switch (node)
         {
             case TypedNodeStatementBlock statement:
@@ -247,8 +260,10 @@ public class Generator
         writer.WriteLine("{");
         using (writer.Indent())
         {
+            PushScope();
             foreach (var statement in node.Statements)
                 EmitStatement(statement);
+            PopScope();
         }
         writer.WriteLine("}");
     }
@@ -262,20 +277,37 @@ public class Generator
 
     private void EmitStatementReturn(TypedNodeStatementReturn statement)
     {
-        var value = EmitExpression(statement.Value);
-        writer.WriteLine($"return {value};");
+        if (statement.Value != null)
+        {
+            var value = EmitExpression(statement.Value);
+            var variableName = TmpName();
+            writer.WriteLine($"{CType(statement.Value.Type, variableName)} = {value};");
+            EmitCleanupAllScopes();
+            writer.WriteLine($"return {variableName};");
+        }
+        else
+        {
+            EmitCleanupAllScopes();
+            writer.WriteLine($"return;");
+        }
+
+        scopes.Peek().Unreachable = true;
     }
 
     private void EmitStatementVariableDeclaration(TypedNodeStatementVariableDeclaration statement)
     {
         var value = EmitExpression(statement.Value);
+        EmitCopyConstructor(value, statement.Value.Type);
         writer.WriteLine($"{CType(statement.Type, statement.Name.Ident)} = {value};");
+        scopes.Peek().Locals.Add((statement.Name.Ident, statement.Type));
     }
 
     private void EmitStatementAssignment(TypedNodeStatementAssignment statement)
     {
         var target = EmitExpression(statement.Target);
+        EmitCopyDestructor(target, statement.Target.Type);
         var value = EmitExpression(statement.Value);
+        EmitCopyConstructor(value, statement.Value.Type);
         writer.WriteLine($"{target} = {value};");
     }
 
@@ -286,7 +318,9 @@ public class Generator
         writer.WriteLine("{");
         using (writer.Indent())
         {
+            PushScope();
             EmitStatement(statement.ThenBlock);
+            PopScope();
         }
         writer.WriteLine("}");
 
@@ -301,7 +335,9 @@ public class Generator
             writer.WriteLine("{");
             using (writer.Indent())
             {
+                PushScope();
                 EmitStatement(statement.ElseBlock);
+                PopScope();
             }
             writer.WriteLine("}");
         }
@@ -314,7 +350,9 @@ public class Generator
         writer.WriteLine("{");
         using (writer.Indent())
         {
+            PushScope();
             EmitStatement(statement.Body);
+            PopScope();
         }
         writer.WriteLine("}");
     }
@@ -342,8 +380,10 @@ public class Generator
                 writer.WriteLine("{");
                 using (writer.Indent())
                 {
+                    PushScope();
                     writer.WriteLine($"{CType(variantInfo.Type, @case.VariableName.Ident)} = {target}.{@case.Variant.Ident};");
                     EmitStatement(@case.Body);
+                    PopScope();
                     writer.WriteLine("break;");
                 }
                 writer.WriteLine("}");
@@ -360,7 +400,7 @@ public class Generator
             TypedNodeExpressionUnary expression => EmitExpressionUnary(expression),
             TypedNodeExpressionBoolLiteral expression => expression.Value.Value ? "true" : "false",
             TypedNodeExpressionIntLiteral expression => expression.Value.Value.ToString(),
-            TypedNodeExpressionStringLiteral expression => $"({CType(expression.Type)}){{ \"{expression.Value.Value}\", {expression.Value.Value.Length} }}",
+            TypedNodeExpressionStringLiteral expression => EmitExpressionStringLiteral(expression),
             TypedNodeExpressionStructLiteral expression => EmitExpressionStructLiteral(expression),
             TypedNodeExpressionEnumLiteral expression => EmitExpressionEnumLiteral(expression),
             TypedNodeExpressionStructMemberAccess expression => EmitExpressionMemberAccess(expression),
@@ -411,6 +451,21 @@ public class Generator
         };
     }
 
+    private string EmitExpressionStringLiteral(TypedNodeExpressionStringLiteral expression)
+    {
+        var name = TmpName();
+        scopes.Peek().Locals.Add((name, expression.Type));
+
+        var variable = CType(expression.Type, name);
+
+        writer.WriteLine($"{variable} = malloc(sizeof(struct {NameMangler.Mangle("core", "string", expression.Type)}));");
+        writer.WriteLine($"{name}->data = \"{expression.Value.Value}\";");
+        writer.WriteLine($"{name}->length = {expression.Value.Value.Length};");
+        writer.WriteLine($"{name}->ref = 1;");
+
+        return name;
+    }
+
     private string EmitExpressionStructLiteral(TypedNodeExpressionStructLiteral expression)
     {
         var initializerValues = new Dictionary<string, string>();
@@ -450,13 +505,13 @@ public class Generator
     private string EmitExpressionStringLength(TypedNodeExpressionStringLength expression)
     {
         var target = EmitExpression(expression.Target);
-        return $"{target}.length";
+        return $"{target}->length";
     }
 
     private string EmitExpressionStringPointer(TypedNodeExpressionStringPointer expression)
     {
         var target = EmitExpression(expression.Target);
-        return $"{target}.data";
+        return $"{target}->data";
     }
 
     private string EmitNodeExpressionGlobalIdent(TypedNodeExpressionGlobalIdent expression)
@@ -489,7 +544,7 @@ public class Generator
             NubTypeSInt type => $"int{type.Width}_t" + (varName != null ? $" {varName}" : ""),
             NubTypeUInt type => $"uint{type.Width}_t" + (varName != null ? $" {varName}" : ""),
             NubTypePointer type => CType(type.To) + (varName != null ? $" *{varName}" : "*"),
-            NubTypeString type => $"struct {NameMangler.Mangle("core", "string", type)}" + (varName != null ? $" {varName}" : ""),
+            NubTypeString type => $"struct {NameMangler.Mangle("core", "string", type)}" + (varName != null ? $" *{varName}" : "*"),
             NubTypeFunc type => $"{CType(type.ReturnType)} (*{varName})({string.Join(", ", type.Parameters.Select(p => CType(p)))})",
             _ => throw new ArgumentOutOfRangeException(nameof(node), node, null)
         };
@@ -499,6 +554,77 @@ public class Generator
     {
         return $"struct {NameMangler.Mangle("anonymous", "struct", type)}{(varName != null ? $" {varName}" : "")}";
     }
+
+    private string TmpName()
+    {
+        return $"_tmp{tmpNameIndex++}";
+    }
+
+    private void EmitCleanupAllScopes()
+    {
+        foreach (var scope in scopes.Reverse())
+        {
+            for (int i = scope.Locals.Count - 1; i >= 0; i--)
+            {
+                var (name, type) = scope.Locals[i];
+                EmitCopyDestructor(name, type);
+            }
+        }
+    }
+
+    private void EmitCleanupCurrentScope(Scope scope)
+    {
+        for (int i = scope.Locals.Count - 1; i >= 0; i--)
+        {
+            var (name, type) = scope.Locals[i];
+            EmitCopyDestructor(name, type);
+        }
+    }
+
+    private void EmitCopyConstructor(string value, NubType type)
+    {
+        switch (type)
+        {
+            case NubTypeString:
+                writer.WriteLine($"{value}->ref += 1;");
+                break;
+        }
+    }
+
+    private void EmitCopyDestructor(string value, NubType type)
+    {
+        switch (type)
+        {
+            case NubTypeString:
+                writer.WriteLine($"{value}->ref -= 1;");
+                writer.WriteLine($"if ({value}->ref == 0)");
+                writer.WriteLine("{");
+                using (writer.Indent())
+                {
+                    writer.WriteLine($"free({value});");
+                }
+                writer.WriteLine("}");
+                break;
+        }
+    }
+
+    private void PushScope()
+    {
+        scopes.Push(new Scope());
+    }
+
+    private void PopScope()
+    {
+        var scope = scopes.Pop();
+        if (!scope.Unreachable)
+            EmitCleanupCurrentScope(scope);
+    }
+
+    private class Scope
+    {
+        public List<(string Name, NubType Type)> Locals { get; } = [];
+        public bool Unreachable { get; set; }
+    }
 }
 
 internal class IndentedTextWriter

compiler/Parser.cs

@@ -195,7 +195,17 @@ public class Parser
 
         if (TryExpectKeyword(Keyword.Return))
         {
-            var value = ParseExpression();
+            NodeExpression? value = null;
+
+            if (Peek() is TokenIdent token && token.Ident == "void")
+            {
+                Next();
+            }
+            else
+            {
+                value = ParseExpression();
+            }
+
             return new NodeStatementReturn(TokensFrom(startIndex), value);
         }
 
@@ -792,9 +802,9 @@ public class NodeStatementExpression(List<Token> tokens, NodeExpression expressi
     public NodeExpression Expression { get; } = expression;
 }
 
-public class NodeStatementReturn(List<Token> tokens, NodeExpression value) : NodeStatement(tokens)
+public class NodeStatementReturn(List<Token> tokens, NodeExpression? value) : NodeStatement(tokens)
 {
-    public NodeExpression Value { get; } = value;
+    public NodeExpression? Value { get; } = value;
 }
 
 public class NodeStatementVariableDeclaration(List<Token> tokens, TokenIdent name, NodeType? type, NodeExpression value) : NodeStatement(tokens)

compiler/Program.cs

@@ -128,7 +128,32 @@ else
     Directory.CreateDirectory(".build");
 }
 
-var output = Generator.Emit(functions, moduleGraph, compileLib);
+string? entryPoint = null;
+
+if (!compileLib)
+{
+    if (!moduleGraph.TryResolveIdentifier("main", "main", true, out var info) || info.Type is not NubTypeFunc entryPointType)
+    {
+        DiagnosticFormatter.Print(Diagnostic.Error("func main::main(): i32 is not defined. If you wanted to compile as a library, specify --type=lib").Build(), Console.Error);
+        return 1;
+    }
+
+    if (!entryPointType.ReturnType.IsAssignableTo(NubTypeSInt.Get(32)))
+    {
+        DiagnosticFormatter.Print(Diagnostic.Error($"Entrypoint must return an i32 (currently '{entryPointType.ReturnType}')").Build(), Console.Error);
+        return 1;
+    }
+
+    if (entryPointType.Parameters.Any())
+    {
+        DiagnosticFormatter.Print(Diagnostic.Error($"Entrypoint must not take any parameters").Build(), Console.Error);
+        return 1;
+    }
+
+    entryPoint = info.MangledName;
+}
+
+var output = Generator.Emit(functions, moduleGraph, entryPoint);
 File.WriteAllText(".build/out.c", output);
 
 if (compileLib)

compiler/TypeChecker.cs

@@ -145,11 +145,21 @@ public class TypeChecker
 
     private TypedNodeStatementReturn CheckStatementReturn(NodeStatementReturn statement)
     {
-        var value = CheckExpression(statement.Value, functionReturnType);
-        if (!value.Type.IsAssignableTo(functionReturnType))
-            throw BasicError($"Type of returned value ({value.Type}) is not assignable to the return type of the function ({functionReturnType})", value);
+        if (statement.Value == null)
+        {
+            if (functionReturnType is not NubTypeVoid)
+                throw BasicError($"Missing return value. Expected '{functionReturnType}'", statement);
 
-        return new TypedNodeStatementReturn(statement.Tokens, value);
+            return new TypedNodeStatementReturn(statement.Tokens, null);
+        }
+        else
+        {
+            var value = CheckExpression(statement.Value, functionReturnType);
+            if (!value.Type.IsAssignableTo(functionReturnType))
+                throw BasicError($"Type of returned value ({value.Type}) is not assignable to the return type of the function ({functionReturnType})", value);
+
+            return new TypedNodeStatementReturn(statement.Tokens, value);
+        }
     }
 
     private TypedNodeStatementVariableDeclaration CheckStatementVariableDeclaration(NodeStatementVariableDeclaration statement)
@@ -796,9 +806,9 @@ public class TypedNodeStatementFuncCall(List<Token> tokens, TypedNodeExpression
     public List<TypedNodeExpression> Parameters { get; } = parameters;
 }
 
-public class TypedNodeStatementReturn(List<Token> tokens, TypedNodeExpression value) : TypedNodeStatement(tokens)
+public class TypedNodeStatementReturn(List<Token> tokens, TypedNodeExpression? value) : TypedNodeStatement(tokens)
 {
-    public TypedNodeExpression Value { get; } = value;
+    public TypedNodeExpression? Value { get; } = value;
 }
 
 public class TypedNodeStatementVariableDeclaration(List<Token> tokens, TokenIdent name, NubType type, TypedNodeExpression value) : TypedNodeStatement(tokens)

examples/build.sh

@@ -1,7 +1,7 @@
 set -e
 
 pushd core
-dotnet run --project ../../compiler mem.nub print.nub --type=lib
+dotnet run --project ../../compiler print.nub --type=lib
 popd
 
 pushd program

examples/core/mem.nub

@@ -1,4 +0,0 @@
-module core
-
-export extern func malloc(size: u64): ^void
-export extern func free(size: ^void)

examples/program/main.nub

@@ -1,7 +1,6 @@
 module main
 
-func main(): i32
-{
-    core::print("Hello, world")
+func main(): i32 {
+    core::print("Your mom")
     return 0
-}
+}