Refs working?

compiler/NubLang/Generation/Generator.cs

@@ -9,10 +9,11 @@ public class Generator
 {
     private readonly CompilationUnit _compilationUnit;
     private readonly IndentedTextWriter _writer;
-    private readonly Stack<List<DeferNode>> _deferStack = [];
+    private readonly Stack<Scope> _scopes = [];
-    private readonly Stack<List<string>> _refCleanupStack = [];
     private int _tmpIndex;
 
+    private Scope Scope => _scopes.Peek();
+
     public Generator(CompilationUnit compilationUnit)
     {
         _compilationUnit = compilationUnit;
@@ -103,7 +104,18 @@ public class Generator
             _writer.WriteLine("{");
             using (_writer.Indent())
             {
-                EmitBlock(funcNode.Body);
+                using (BeginScope())
+                {
+                    foreach (var parameter in funcNode.Prototype.Parameters)
+                    {
+                        if (parameter.Type is NubRefType)
+                        {
+                            Scope.Defer(() => _writer.WriteLine($"rc_release({parameter.NameToken.Value});"));
+                        }
+                    }
+
+                    EmitBlock(funcNode.Body);
+                }
             }
 
             _writer.WriteLine("}");
@@ -125,7 +137,10 @@ public class Generator
                 _writer.WriteLine("{");
                 using (_writer.Indent())
                 {
-                    EmitBlock(blockNode);
+                    using (BeginScope())
+                    {
+                        EmitBlock(blockNode);
+                    }
                 }
 
                 _writer.WriteLine("}");
@@ -137,7 +152,7 @@ public class Generator
                 EmitContinue(continueNode);
                 break;
             case DeferNode deferNode:
-                EmitDefer(deferNode);
+                Scope.Defer(() => EmitStatement(deferNode.Statement));
                 break;
             case ForConstArrayNode forConstArrayNode:
                 EmitForConstArray(forConstArrayNode);
@@ -177,6 +192,13 @@ public class Generator
     {
         var target = EmitExpression(assignmentNode.Target);
         var value = EmitExpression(assignmentNode.Value);
+
+        if (assignmentNode.Target.Type is NubRefType)
+        {
+            _writer.WriteLine($"rc_retain({value});");
+            _writer.WriteLine($"rc_release({target});");
+        }
+
         _writer.WriteLine($"{target} = {value};");
     }
 
@@ -192,11 +214,6 @@ public class Generator
         _writer.WriteLine("continue;");
     }
 
-    private void EmitDefer(DeferNode deferNode)
-    {
-        _deferStack.Peek().Add(deferNode);
-    }
-
     private void EmitForSlice(ForSliceNode forSliceNode)
     {
         var targetType = (NubSliceType)forSliceNode.Target.Type;
@@ -207,8 +224,11 @@ public class Generator
         _writer.WriteLine("{");
         using (_writer.Indent())
         {
-            _writer.WriteLine($"{CType.Create(targetType.ElementType, forSliceNode.ElementNameToken.Value)} = (({CType.Create(targetType.ElementType)}*){target}.data)[{indexName}];");
+            using (BeginScope())
-            EmitBlock(forSliceNode.Body);
+            {
+                _writer.WriteLine($"{CType.Create(targetType.ElementType, forSliceNode.ElementNameToken.Value)} = (({CType.Create(targetType.ElementType)}*){target}.data)[{indexName}];");
+                EmitBlock(forSliceNode.Body);
+            }
         }
 
         _writer.WriteLine("}");
@@ -224,8 +244,11 @@ public class Generator
         _writer.WriteLine("{");
         using (_writer.Indent())
         {
-            _writer.WriteLine($"{CType.Create(targetType.ElementType, forConstArrayNode.ElementNameToken.Value)} = {target}[{indexName}];");
+            using (BeginScope())
-            EmitBlock(forConstArrayNode.Body);
+            {
+                _writer.WriteLine($"{CType.Create(targetType.ElementType, forConstArrayNode.ElementNameToken.Value)} = {target}[{indexName}];");
+                EmitBlock(forConstArrayNode.Body);
+            }
         }
 
         _writer.WriteLine("}");
@@ -238,7 +261,10 @@ public class Generator
         _writer.WriteLine("{");
         using (_writer.Indent())
         {
-            EmitBlock(ifNode.Body);
+            using (BeginScope())
+            {
+                EmitBlock(ifNode.Body);
+            }
         }
 
         _writer.WriteLine("}");
@@ -251,7 +277,10 @@ public class Generator
                 _writer.WriteLine("{");
                 using (_writer.Indent())
                 {
-                    EmitBlock(elseNode);
+                    using (BeginScope())
+                    {
+                        EmitBlock(elseNode);
+                    }
                 }
 
                 _writer.WriteLine("}");
@@ -263,39 +292,20 @@ public class Generator
     {
         if (returnNode.Value == null)
         {
-            var blockDefers = _deferStack.Peek();
+            EmitScopeCleanup();
-            for (var i = blockDefers.Count - 1; i >= 0; i--)
-            {
-                EmitStatement(blockDefers[i].Statement);
-            }
-
-            var refCleanups = _refCleanupStack.Peek();
-            foreach (var refCleanup in refCleanups)
-            {
-                _writer.WriteLine($"rc_release({refCleanup});");
-            }
-
             _writer.WriteLine("return;");
         }
         else
         {
             var returnValue = EmitExpression(returnNode.Value);
-
             var tmp = NewTmp();
             _writer.WriteLine($"{CType.Create(returnNode.Value.Type, tmp)} = {returnValue};");
-
+            if (returnNode.Value.Type is NubRefType)
-            var blockDefers = _deferStack.Peek();
-            for (var i = blockDefers.Count - 1; i >= 0; i--)
             {
-                EmitStatement(blockDefers[i].Statement);
+                _writer.WriteLine($"rc_retain({tmp});");
-            }
-
-            var refCleanups = _refCleanupStack.Peek();
-            foreach (var refCleanup in refCleanups)
-            {
-                _writer.WriteLine($"rc_release({refCleanup});");
             }
 
+            EmitScopeCleanup();
             EmitLine(returnNode.Tokens.FirstOrDefault());
             _writer.WriteLine($"return {tmp};");
         }
@@ -313,6 +323,10 @@ public class Generator
         {
             var value = EmitExpression(variableDeclarationNode.Assignment);
             _writer.WriteLine($"{CType.Create(variableDeclarationNode.Type, variableDeclarationNode.NameToken.Value)} = {value};");
+            if (variableDeclarationNode.Type is NubRefType)
+            {
+                _writer.WriteLine($"rc_retain({variableDeclarationNode.NameToken.Value});");
+            }
         }
         else
         {
@@ -327,7 +341,10 @@ public class Generator
         _writer.WriteLine("{");
         using (_writer.Indent())
         {
-            EmitBlock(whileNode.Body);
+            using (BeginScope())
+            {
+                EmitBlock(whileNode.Body);
+            }
         }
 
         _writer.WriteLine("}");
@@ -492,7 +509,19 @@ public class Generator
     private string EmitFuncCall(FuncCallNode funcCallNode)
     {
         var name = EmitExpression(funcCallNode.Expression);
-        var parameterNames = funcCallNode.Parameters.Select(EmitExpression).ToList();
+
+        var parameterNames = new List<string>();
+        foreach (var parameter in funcCallNode.Parameters)
+        {
+            var result = EmitExpression(parameter);
+            if (parameter.Type is NubRefType)
+            {
+                _writer.WriteLine($"rc_retain({result});");
+            }
+
+            parameterNames.Add(result);
+        }
+
         return $"{name}({string.Join(", ", parameterNames)})";
     }
 
@@ -529,7 +558,7 @@ public class Generator
 
         _writer.WriteLine($"*{tmp} = ({CType.Create(structType)}){{{initString}}};");
 
-        _refCleanupStack.Peek().Add(tmp);
+        Scope.Defer(() => _writer.WriteLine($"rc_release({tmp});"));
 
         return tmp;
     }
@@ -625,24 +654,45 @@ public class Generator
 
     private void EmitBlock(BlockNode blockNode)
     {
-        _deferStack.Push([]);
-        _refCleanupStack.Push([]);
-
         foreach (var statementNode in blockNode.Statements)
         {
             EmitStatement(statementNode);
         }
+    }
 
-        var blockDefers = _deferStack.Pop();
+    private void EmitScopeCleanup()
-        for (var i = blockDefers.Count - 1; i >= 0; i--)
+    {
+        var deferredStack = Scope.GetDeferred();
+        while (deferredStack.TryPop(out var deferred))
         {
-            EmitStatement(blockDefers[i].Statement);
+            deferred.Invoke();
         }
+    }
 
-        var refCleanups = _refCleanupStack.Pop();
+    private ScopeDisposer BeginScope()
-        foreach (var refCleanup in refCleanups)
+    {
+        _scopes.Push(new Scope());
+        return new ScopeDisposer(this);
+    }
+
+    private sealed class ScopeDisposer(Generator owner) : IDisposable
+    {
+        public void Dispose()
         {
-            _writer.WriteLine($"rc_release({refCleanup});");
+            owner.EmitScopeCleanup();
+            owner._scopes.Pop();
         }
     }
 }
+
+public class Scope
+{
+    private readonly List<Action> _deferred = [];
+
+    public void Defer(Action action)
+    {
+        _deferred.Add(action);
+    }
+
+    public Stack<Action> GetDeferred() => new(_deferred);
+}

examples/playgroud/main.nub

@@ -16,5 +16,12 @@ extern "main" func main(argc: i64, argv: [?]^i8): i64
     }
 
     puts(x^.name)
+
+    test(x)
+
     return 0
 }
+
+func test(x: &Human)
+{
+}

runtime/ref.c

@@ -12,7 +12,7 @@ void *rc_alloc(size_t size, void (*destructor)(void *self))
         exit(69);
     }
 
-    header->ref_count = 1;
+    header->ref_count = 0;
     header->destructor = destructor;
 
     return (void *)(header + 1);
@@ -29,7 +29,7 @@ void rc_release(void *obj)
 {
     ref_header *header = ((ref_header *)obj) - 1;
     printf("rc_release\n");
-    if (--header->ref_count == 0)
+    if (--header->ref_count <= 0)
     {
         if (header->destructor)
         {