Use enum instead of symbol for Atomic primitives (#11583)

spec/compiler/codegen/primitives_spec.cr

@@ -310,6 +310,61 @@ describe "Code gen: primitives" do
     {% end %}
   end
 
+  describe "atomicrmw" do
+    it "codegens atomicrmw with enums" do
+      run(<<-CR).to_i.should eq(3)
+        enum RMWBinOp
+          Add = 1
+        end
+
+        enum Ordering
+          SequentiallyConsistent = 7
+        end
+
+        @[Primitive(:atomicrmw)]
+        def atomicrmw(op : RMWBinOp, ptr : Int32*, val : Int32, ordering : Ordering, singlethread : Bool) : Int32
+        end
+
+        x = 1
+        atomicrmw(:add, pointerof(x), 2, :sequentially_consistent, false)
+        x
+        CR
+    end
+
+    it "codegens atomicrmw with enums" do
+      run(<<-CR).to_i.should eq(3)
+        enum RMWBinOp
+          Add = 1
+        end
+
+        enum Ordering
+          SequentiallyConsistent = 7
+        end
+
+        @[Primitive(:atomicrmw)]
+        def atomicrmw(op : RMWBinOp, ptr : Int32*, val : Int32, ordering : Ordering, singlethread : Bool) : Int32
+        end
+
+        x = 1
+        atomicrmw(RMWBinOp::Add, pointerof(x), 2, Ordering::SequentiallyConsistent, false)
+        x
+        CR
+    end
+
+    # TODO: remove once support for 1.4 is dropped
+    it "codegens atomicrmw with symbols" do
+      run(<<-CR).to_i.should eq(3)
+        @[Primitive(:atomicrmw)]
+        def atomicrmw(op : Symbol, ptr : Int32*, val : Int32, ordering : Symbol, singlethread : Bool) : Int32
+        end
+
+        x = 1
+        atomicrmw(:add, pointerof(x), 2, :sequentially_consistent, false)
+        x
+        CR
+    end
+  end
+
   it "allows @[Primitive] on method that has body" do
     run(%(
       module Moo

src/atomic.cr

@@ -1,3 +1,5 @@
+require "llvm/enums/atomic"
+
 # A value that may be updated atomically.
 #
 # Only primitive integer types, reference types or nilable reference types
@@ -202,23 +204,23 @@ struct Atomic(T)
     # Defines methods that directly map to LLVM instructions related to atomic operations.
 
     @[Primitive(:cmpxchg)]
-    def self.cmpxchg(ptr : T*, cmp : T, new : T, success_ordering : Symbol, failure_ordering : Symbol) : {T, Bool} forall T
+    def self.cmpxchg(ptr : T*, cmp : T, new : T, success_ordering : LLVM::AtomicOrdering, failure_ordering : LLVM::AtomicOrdering) : {T, Bool} forall T
     end
 
     @[Primitive(:atomicrmw)]
-    def self.atomicrmw(op : Symbol, ptr : T*, val : T, ordering : Symbol, singlethread : Bool) : T forall T
+    def self.atomicrmw(op : LLVM::AtomicRMWBinOp, ptr : T*, val : T, ordering : LLVM::AtomicOrdering, singlethread : Bool) : T forall T
     end
 
     @[Primitive(:fence)]
-    def self.fence(ordering : Symbol, singlethread : Bool) : Nil
+    def self.fence(ordering : LLVM::AtomicOrdering, singlethread : Bool) : Nil
     end
 
     @[Primitive(:load_atomic)]
-    def self.load(ptr : T*, ordering : Symbol, volatile : Bool) : T forall T
+    def self.load(ptr : T*, ordering : LLVM::AtomicOrdering, volatile : Bool) : T forall T
     end
 
     @[Primitive(:store_atomic)]
-    def self.store(ptr : T*, value : T, ordering : Symbol, volatile : Bool) : Nil forall T
+    def self.store(ptr : T*, value : T, ordering : LLVM::AtomicOrdering, volatile : Bool) : Nil forall T
     end
   end
 end

src/compiler/crystal/codegen/primitives.cr

@@ -1162,10 +1162,11 @@ class Crystal::CodeGenVisitor
 
   def codegen_primitive_cmpxchg(call, node, target_def, call_args)
     call = check_atomic_call(call, target_def)
-    success_ordering = atomic_ordering_from_symbol_literal(call.args[-2])
-    failure_ordering = atomic_ordering_from_symbol_literal(call.args[-1])
+    ptr, cmp, new, success_ordering, failure_ordering = call_args
+
+    success_ordering = atomic_ordering_get_const(call.args[-2], success_ordering)
+    failure_ordering = atomic_ordering_get_const(call.args[-1], failure_ordering)
 
-    ptr, cmp, new, _, _ = call_args
     value = builder.cmpxchg(ptr, cmp, new, success_ordering, failure_ordering)
     value_ptr = alloca llvm_type(node.type)
     store extract_value(value, 0), gep(value_ptr, 0, 0)
@@ -1175,17 +1176,20 @@ class Crystal::CodeGenVisitor
 
   def codegen_primitive_atomicrmw(call, node, target_def, call_args)
     call = check_atomic_call(call, target_def)
-    op = atomicrwm_bin_op_from_symbol_literal(call.args[0])
-    ordering = atomic_ordering_from_symbol_literal(call.args[-2])
+    op, ptr, val, ordering, _ = call_args
+
+    op = atomicrwm_bin_op_get_const(call.args[0], op)
+    ordering = atomic_ordering_get_const(call.args[-2], ordering)
     singlethread = bool_from_bool_literal(call.args[-1])
 
-    _, ptr, val, _, _ = call_args
     builder.atomicrmw(op, ptr, val, ordering, singlethread)
   end
 
   def codegen_primitive_fence(call, node, target_def, call_args)
     call = check_atomic_call(call, target_def)
-    ordering = atomic_ordering_from_symbol_literal(call.args[0])
+    ordering, _ = call_args
+
+    ordering = atomic_ordering_get_const(call.args[0], ordering)
     singlethread = bool_from_bool_literal(call.args[1])
 
     builder.fence(ordering, singlethread)
@@ -1194,10 +1198,10 @@ class Crystal::CodeGenVisitor
 
   def codegen_primitive_load_atomic(call, node, target_def, call_args)
     call = check_atomic_call(call, target_def)
-    ordering = atomic_ordering_from_symbol_literal(call.args[-2])
-    volatile = bool_from_bool_literal(call.args[-1])
+    ptr, ordering, _ = call_args
 
-    ptr, _, _ = call_args
+    ordering = atomic_ordering_get_const(call.args[-2], ordering)
+    volatile = bool_from_bool_literal(call.args[-1])
 
     inst = builder.load(ptr)
     inst.ordering = ordering
@@ -1208,10 +1212,10 @@ class Crystal::CodeGenVisitor
 
   def codegen_primitive_store_atomic(call, node, target_def, call_args)
     call = check_atomic_call(call, target_def)
-    ordering = atomic_ordering_from_symbol_literal(call.args[-2])
-    volatile = bool_from_bool_literal(call.args[-1])
+    ptr, value, ordering, _ = call_args
 
-    ptr, value, _, _ = call_args
+    ordering = atomic_ordering_get_const(call.args[-2], ordering)
+    volatile = bool_from_bool_literal(call.args[-1])
 
     inst = builder.store(value, ptr)
     inst.ordering = ordering
@@ -1246,30 +1250,39 @@ class Crystal::CodeGenVisitor
     end
   end
 
-  def atomic_ordering_from_symbol_literal(node)
-    unless node.is_a?(SymbolLiteral)
-      node.raise "BUG: expected symbol literal"
-    end
+  def atomic_ordering_get_const(node, llvm_arg)
+    node.raise "atomic ordering must be a constant" unless llvm_arg.constant?
 
-    ordering = LLVM::AtomicOrdering.parse?(node.value)
-    unless ordering
-      node.raise "unknown atomic ordering: #{node.value}"
+    if node.type.implements?(@program.enum) && llvm_arg.type.kind.integer? && llvm_arg.type.int_width == 32
+      # any `Int32` enum will do, it is up to `Atomic::Ops` to use appropriate
+      # parameter restrictions so that things don't go bad
+      LLVM::AtomicOrdering.new(llvm_arg.const_int_get_sext_value.to_i32!)
+    elsif node.is_a?(SymbolLiteral)
+      # TODO: remove once support for 1.4 is dropped
+      op = LLVM::AtomicOrdering.parse?(node.value)
+      unless op
+        node.raise "unknown atomic ordering: #{node.value}"
+      end
+      op
+    else
+      node.raise "BUG: unknown atomic ordering: #{node}"
     end
-
-    ordering
   end
 
-  def atomicrwm_bin_op_from_symbol_literal(node)
-    unless node.is_a?(SymbolLiteral)
-      node.raise "BUG: expected symbol literal"
-    end
+  def atomicrwm_bin_op_get_const(node, llvm_arg)
+    node.raise "atomic rwm bin op must be a constant" unless llvm_arg.constant?
 
-    op = LLVM::AtomicRMWBinOp.parse?(node.value)
-    unless op
-      node.raise "unknown atomic rwm bin op: #{node.value}"
+    if node.type.implements?(@program.enum) && llvm_arg.type.kind.integer? && llvm_arg.type.int_width == 32
+      LLVM::AtomicRMWBinOp.new(llvm_arg.const_int_get_sext_value.to_i32!)
+    elsif node.is_a?(SymbolLiteral)
+      op = LLVM::AtomicRMWBinOp.parse?(node.value)
+      unless op
+        node.raise "unknown atomic rwm bin op: #{node.value}"
+      end
+      op
+    else
+      node.raise "BUG: unknown atomic rwm bin op: #{node}"
     end
-
-    op
   end
 
   def bool_from_bool_literal(node)

src/llvm/enums.cr

@@ -361,30 +361,6 @@ module LLVM
     HiUser         = 0xff
   end
 
-  enum AtomicOrdering
-    NotAtomic              = 0
-    Unordered              = 1
-    Monotonic              = 2
-    Acquire                = 4
-    Release                = 5
-    AcquireRelease         = 6
-    SequentiallyConsistent = 7
-  end
-
-  enum AtomicRMWBinOp
-    Xchg
-    Add
-    Sub
-    And
-    Nand
-    Or
-    Xor
-    Max
-    Min
-    UMax
-    UMin
-  end
-
   enum DIFlags : UInt32
     Zero                = 0
     Private             = 1
@@ -487,3 +463,5 @@ module LLVM
     end
   end
 end
+
+require "./enums/*"

src/llvm/enums/atomic.cr

@@ -0,0 +1,27 @@
+module LLVM
+  enum AtomicOrdering
+    NotAtomic              = 0
+    Unordered              = 1
+    Monotonic              = 2
+    Acquire                = 4
+    Release                = 5
+    AcquireRelease         = 6
+    SequentiallyConsistent = 7
+  end
+
+  enum AtomicRMWBinOp
+    Xchg
+    Add
+    Sub
+    And
+    Nand
+    Or
+    Xor
+    Max
+    Min
+    Umax
+    Umin
+    Fadd
+    Fsub
+  end
+end