From d4b0384626ac4adb1bd0b5025e4b1c7493c970c6 Mon Sep 17 00:00:00 2001
From: Patrick Smith <pgwsmith@gmail.com>
Date: Mon, 4 Nov 2024 21:44:01 +1100
Subject: [PATCH] Allow returning atoms using str constant system

---
 lib/orb/constants.ex           |  6 ++++++
 test/string_constants_test.exs | 35 ++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/lib/orb/constants.ex b/lib/orb/constants.ex
index 8d21ee0..a59299e 100644
--- a/lib/orb/constants.ex
+++ b/lib/orb/constants.ex
@@ -128,6 +128,12 @@ defmodule Orb.Constants do
   def expand_if_needed(value) when is_binary(value), do: expand_string!(value)
   def expand_if_needed(value) when is_list(value), do: :lists.map(&expand_if_needed/1, value)
   def expand_if_needed(value) when is_struct(value, Orb.IfElse), do: Orb.IfElse.expand(value)
+  def expand_if_needed(nil), do: nil
+
+  def expand_if_needed(value) when is_atom(value) do
+    str = value |> Atom.to_string() |> expand_string!()
+    str.memory_offset
+  end
 
   # Handles Orb.InstructionSequence or anything with a `body`
   def expand_if_needed(%_{body: _} = struct) do
diff --git a/test/string_constants_test.exs b/test/string_constants_test.exs
index 4970836..cb4be88 100644
--- a/test/string_constants_test.exs
+++ b/test/string_constants_test.exs
@@ -32,6 +32,41 @@ defmodule StringConstantsTest do
            """
   end
 
+  test "works with atoms" do
+    defmodule ABCAtom do
+      use Orb
+
+      defw alphabet(), I32 do
+        :abc
+      end
+
+      defw always_be_closing(), I32 do
+        :abc
+      end
+
+      defw australian_broadcasting_corporation(), Str do
+        "abc"
+      end
+    end
+
+    assert to_wat(ABCAtom) == """
+           (module $ABCAtom
+             (memory (export "memory") 1)
+             (; constants 4 bytes ;)
+             (data (i32.const 255) "abc")
+             (func $alphabet (export "alphabet") (result i32)
+               (i32.const 255)
+             )
+             (func $always_be_closing (export "always_be_closing") (result i32)
+               (i32.const 255)
+             )
+             (func $australian_broadcasting_corporation (export "australian_broadcasting_corporation") (result i32 i32)
+               (i32.const 255) (i32.const 3)
+             )
+           )
+           """
+  end
+
   describe "const/1" do
     test "assigns data" do
       defmodule ConstHTMLTypes do