Fix: file loading handling

src/Lua/CodeAnalysis/Compilation/Scanner.cs

@@ -221,6 +221,24 @@ public void Advance()
         Current = R.TryRead(out var c) ? c : EndOfStream;
     }
 
+    void SkipFirstLineComment()
+    {
+        if (Current != '#')
+        {
+            return;
+        }
+
+        while (!IsNewLine(Current) && Current != EndOfStream)
+        {
+            Advance();
+        }
+
+        if (IsNewLine(Current))
+        {
+            IncrementLineNumber();
+        }
+    }
+
     public void SaveAndAdvance()
     {
         Save(Current);
@@ -806,7 +824,15 @@ public Token Scan()
 
                     return ReadNumber(pos);
                 case 0:
-                    Advance();
+                    if (R.Position == 0)
+                    {
+                        Advance();
+                        SkipFirstLineComment();
+                    }
+                    else
+                    {
+                        Advance();
+                    }
                     pos = R.Position;
                     break;
                 default:

src/Lua/IO/FileSystem.cs

@@ -81,6 +81,11 @@ public ValueTask Rename(string oldName, string newName, CancellationToken cancel
     public ValueTask Remove(string path, CancellationToken cancellationToken)
     {
         path = GetFullPath(path);
+        if (!File.Exists(path))
+        {
+            throw new FileNotFoundException("No such file or directory", path);
+        }
+
         File.Delete(path);
         return default;
     }

src/Lua/IO/LuaStream.cs

@@ -33,6 +33,11 @@ public ValueTask<string> ReadAllAsync(CancellationToken cancellationToken)
     {
         mode.ThrowIfNotReadable();
         reader ??= new();
+        if (count == 0)
+        {
+            return new(reader.IsEndOfStream(innerStream) ? null : string.Empty);
+        }
+
         return new(reader.Read(innerStream, count));
     }
 
@@ -64,13 +69,25 @@ public ValueTask WriteAsync(ReadOnlyMemory<char> buffer, CancellationToken cance
 
         using PooledArray<byte> byteBuffer = new(4096);
         var encoder = Encoding.UTF8.GetEncoder();
-        var totalBytes = encoder.GetByteCount(buffer.Span, true);
-        var remainingBytes = totalBytes;
-        while (0 < remainingBytes)
+        var remainingChars = buffer.Span;
+        var totalBytes = 0;
+        while (!remainingChars.IsEmpty)
         {
-            var byteCount = encoder.GetBytes(buffer.Span, byteBuffer.AsSpan(), false);
-            innerStream.Write(byteBuffer.AsSpan()[..byteCount]);
-            remainingBytes -= byteCount;
+            encoder.Convert(
+                remainingChars,
+                byteBuffer.AsSpan(),
+                flush: true,
+                out var charsUsed,
+                out var bytesUsed,
+                out _);
+
+            if (bytesUsed > 0)
+            {
+                innerStream.Write(byteBuffer.AsSpan()[..bytesUsed]);
+                totalBytes += bytesUsed;
+            }
+
+            remainingChars = remainingChars[charsUsed..];
         }
 
         if (nextFlushSize < (ulong)totalBytes)

src/Lua/Internal/Utf8Reader.cs

@@ -6,11 +6,9 @@ namespace Lua.Internal;
 
 sealed class Utf8Reader
 {
-    [ThreadStatic]
-    static byte[]? scratchBuffer;
+    [ThreadStatic] static byte[]? scratchBuffer;
 
-    [ThreadStatic]
-    internal static bool scratchBufferUsed;
+    [ThreadStatic] internal static bool scratchBufferUsed;
 
     readonly byte[] buffer;
     int bufPos, bufLen;
@@ -235,7 +233,7 @@ public byte ReadByte(Stream stream)
                 }
             }
 
-            if (!dataRead || len != charCount)
+            if (!dataRead)
             {
                 return null;
             }
@@ -415,6 +413,11 @@ public void Clear()
         }
     }
 
+    public bool IsEndOfStream(Stream stream)
+    {
+        return PeekByte(stream) < 0;
+    }
+
     int PeekByte(Stream stream, int offset = 0)
     {
         // Ensure we have enough data in buffer

src/Lua/Runtime/LuaClosure.cs

@@ -62,6 +62,22 @@ internal void SetUpValue(int index, LuaValue value)
         upValues[index].SetValue(value);
     }
 
+    internal void SetEnvironment(LuaValue environment)
+    {
+        if (Proto.UpValues.Length == 0 || Proto.UpValues[0].Name != "_ENV")
+        {
+            return;
+        }
+
+        if (upValues.Length == 0)
+        {
+            upValues.Add(UpValue.Closed(environment));
+            return;
+        }
+
+        upValues[0] = UpValue.Closed(environment);
+    }
+
     static UpValue GetUpValueFromDescription(LuaGlobalState globalState, LuaState state, UpValueDesc description, int baseIndex = 0)
     {
         if (description.IsLocal)

src/Lua/Standard/BasicLibrary.cs

@@ -9,6 +9,7 @@ namespace Lua.Standard;
 public sealed class BasicLibrary
 {
     public static readonly BasicLibrary Instance = new();
+
     static readonly HashSet<string> KnownCollectGarbageOptions = new(StringComparer.Ordinal)
     {
         "collect",
@@ -182,14 +183,21 @@ public async ValueTask<int> LoadFile(LuaFunctionExecutionContext context, Cancel
         var mode = context.HasArgument(1)
             ? context.GetArgument<string>(1)
             : "bt";
-        var arg2 = context.HasArgument(2)
-            ? context.GetArgument<LuaTable>(2)
-            : null;
+        var hasEnvironment = context.ArgumentCount > 2;
+        var environment = hasEnvironment
+            ? context.GetArgument(2)
+            : LuaValue.Nil;
 
         // do not use LuaState.DoFileAsync as it uses the newExecutionContext
         try
         {
-            return context.Return(await context.State.LoadFileAsync(arg0, mode, arg2, cancellationToken));
+            var closure = await context.State.LoadFileAsync(arg0, mode, null, cancellationToken);
+            if (hasEnvironment)
+            {
+                closure.SetEnvironment(environment);
+            }
+
+            return context.Return(closure);
         }
         catch (Exception ex)
         {
@@ -199,7 +207,6 @@ public async ValueTask<int> LoadFile(LuaFunctionExecutionContext context, Cancel
 
     public ValueTask<int> Load(LuaFunctionExecutionContext context, CancellationToken cancellationToken)
     {
-        // Lua-CSharp does not support binary chunks, the mode argument is ignored.
         var arg0 = context.GetArgument(0);
 
         var name = context.HasArgument(1)
@@ -210,16 +217,28 @@ public ValueTask<int> Load(LuaFunctionExecutionContext context, CancellationToke
             ? context.GetArgument<string>(2)
             : "bt";
 
-        var arg3 = context.HasArgument(3)
-            ? context.GetArgument<LuaTable>(3)
-            : null;
+        var hasEnvironment = context.ArgumentCount > 3;
+        var environment = hasEnvironment
+            ? context.GetArgument(3)
+            : LuaValue.Nil;
 
         // do not use LuaState.DoFileAsync as it uses the newExecutionContext
         try
         {
             if (arg0.TryRead<string>(out var str))
             {
-                return new(context.Return(context.State.Load(str, name ?? str, arg3)));
+                if (!mode.Contains('t'))
+                {
+                    throw new Exception("attempt to load a text chunk (mode is 'b')");
+                }
+
+                var closure = context.State.Load(str, name ?? str, null);
+                if (hasEnvironment)
+                {
+                    closure.SetEnvironment(environment);
+                }
+
+                return new(context.Return(closure));
             }
             else if (arg0.TryRead<LuaFunction>(out var function))
             {

src/Lua/Standard/FileHandle.cs

@@ -119,7 +119,7 @@ public void SetVBuf(string mode, int size)
 
     public async ValueTask Close(CancellationToken cancellationToken)
     {
-        if (!stream.IsOpen)
+        if (stream == null || !stream.IsOpen)
             throw new ObjectDisposedException(nameof(FileHandle));
 
         await stream.CloseAsync(cancellationToken);

src/Lua/Standard/IOLibrary.cs

@@ -86,34 +86,43 @@ public async ValueTask<int> Input(LuaFunctionExecutionContext context, Cancellat
 
     public async ValueTask<int> Lines(LuaFunctionExecutionContext context, CancellationToken cancellationToken)
     {
-        if (context.ArgumentCount == 0)
+        if (context.ArgumentCount == 0 || context.Arguments[0].Type is LuaValueType.Nil)
         {
             var file = context.GlobalState.Registry["_IO_input"].Read<FileHandle>();
-            return context.Return(new CSharpClosure("iterator", [new(file)], static async (context, cancellationToken) =>
+            LuaValue[] upValues = new LuaValue[context.ArgumentCount == 0 ? 1 : context.ArgumentCount];
+            upValues[0] = new(file);
+            if (context.ArgumentCount > 1)
             {
-                var file = context.GetCsClosure()!.UpValues[0].Read<FileHandle>();
-                context.Return();
-                var resultCount = await IOHelper.ReadAsync(context.State, file, "io.lines", 0, Memory<LuaValue>.Empty, true, cancellationToken);
-                if (resultCount > 0 && context.State.Stack.Get(context.ReturnFrameBase).Type is LuaValueType.Nil)
-                {
-                    await file.Close(cancellationToken);
-                }
+                context.Arguments.Slice(1, context.ArgumentCount - 1).CopyTo(upValues.AsSpan(1));
+            }
 
+            return context.Return(new CSharpClosure("iterator", upValues, static async (context, cancellationToken) =>
+            {
+                var upValues = context.GetCsClosure()!.UpValues.AsMemory();
+                var file = upValues.Span[0].Read<FileHandle>();
+                context.Return();
+                var resultCount = await IOHelper.ReadAsync(context.State, file, "io.lines", 0, upValues[1..], true, cancellationToken);
                 return resultCount;
             }));
         }
         else
         {
             var fileName = context.GetArgument<string>(0);
+            LuaValue[] formats = context.ArgumentCount > 1
+                ? context.Arguments[1..context.ArgumentCount].ToArray()
+                : [];
             var stack = context.State.Stack;
             context.Return();
 
             await IOHelper.Open(context.State, fileName, "r", true, cancellationToken);
 
             var file = stack.Get(context.ReturnFrameBase).Read<FileHandle>();
-            var upValues = new LuaValue[context.Arguments.Length];
+            var upValues = new LuaValue[formats.Length + 1];
             upValues[0] = new(file);
-            context.Arguments[1..].CopyTo(upValues[1..]);
+            if (formats.Length > 0)
+            {
+                formats.CopyTo(upValues.AsSpan(1));
+            }
 
             return context.Return(new CSharpClosure("iterator", upValues, static async (context, cancellationToken) =>
             {
@@ -154,6 +163,9 @@ public async ValueTask<int> Open(LuaFunctionExecutionContext context, Cancellati
     public async ValueTask<int> Output(LuaFunctionExecutionContext context, CancellationToken cancellationToken)
     {
         var io = context.GlobalState.Registry;
+        var previousOutput = io["_IO_output"].TryRead<FileHandle>(out var currentOutput)
+            ? currentOutput
+            : null;
 
         if (context.ArgumentCount == 0 || context.Arguments[0].Type is LuaValueType.Nil)
         {
@@ -163,11 +175,21 @@ public async ValueTask<int> Output(LuaFunctionExecutionContext context, Cancella
         var arg = context.Arguments[0];
         if (arg.TryRead<FileHandle>(out var file))
         {
+            if (!ReferenceEquals(previousOutput, file) && previousOutput is { IsOpen: true })
+            {
+                await previousOutput.FlushAsync(cancellationToken);
+            }
+
             io["_IO_output"] = new(file);
             return context.Return(new LuaValue(file));
         }
         else
         {
+            if (previousOutput is { IsOpen: true })
+            {
+                await previousOutput.FlushAsync(cancellationToken);
+            }
+
             var stream = await context.GlobalState.Platform.FileSystem.Open(arg.ToString(), LuaFileOpenMode.Write, cancellationToken);
             FileHandle handle = new(stream);
             io["_IO_output"] = new(handle);

src/Lua/Standard/Internal/IOHelper.cs

@@ -37,6 +37,11 @@ public static async ValueTask<int> Open(LuaState state, string fileName, string
     // TODO: optimize (use IBuffertWrite<byte>, async)
     public static async ValueTask<int> WriteAsync(FileHandle file, string name, LuaFunctionExecutionContext context, CancellationToken cancellationToken)
     {
+        if (!file.IsOpen)
+        {
+            throw new LuaRuntimeException(context.State, "attempt to use a closed file");
+        }
+
         try
         {
             for (var i = 0; i < context.ArgumentCount; i++)
@@ -78,6 +83,11 @@ public static async ValueTask<int> WriteAsync(FileHandle file, string name, LuaF
 
     public static async ValueTask<int> ReadAsync(LuaState state, FileHandle file, string name, int startArgumentIndex, ReadOnlyMemory<LuaValue> formats, bool throwError, CancellationToken cancellationToken)
     {
+        if (!file.IsOpen)
+        {
+            throw new LuaRuntimeException(state, "attempt to use a closed file");
+        }
+
         if (formats.Length == 0)
         {
             formats = defaultReadFormat;
@@ -110,7 +120,7 @@ public static async ValueTask<int> ReadAsync(LuaState state, FileHandle file, st
                         case "L":
                         case "*L":
                             var text = await file.ReadLineAsync(true, cancellationToken);
-                            stack.Push(text == null ? LuaValue.Nil : text + Environment.NewLine);
+                            stack.Push(text == null ? LuaValue.Nil : text);
                             break;
                     }
                 }