Enum: automatically generate a none? method for a Flags enum

crystal-lang · Nov 22, 2016 · 202e771 · 202e771
1 parent eab0094
commit 202e771
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/spec/compiler/codegen/enum_spec.cr b/spec/compiler/codegen/enum_spec.cr
@@ -303,4 +303,19 @@ describe "Code gen: enum" do
       Foo::A.value
       )).to_i.should eq(30)
   end
+
+  it "adds a none? method to flags enum" do
+    run(%(
+      @[Flags]
+      enum Foo
+        A
+        B
+      end
+
+      x = 0
+      x += 1 if Foo::None.none?
+      x += 2 if Foo::A.none?
+      x
+      )).to_i.should eq(1)
+  end
 end
diff --git a/src/compiler/crystal/semantic/top_level_visitor.cr b/src/compiler/crystal/semantic/top_level_visitor.cr
@@ -477,6 +477,8 @@ class Crystal::TopLevelVisitor < Crystal::SemanticVisitor
           none = NumberLiteral.new(0, enum_base_type.kind)
           none.type = enum_type
           enum_type.add_constant Arg.new("None", default_value: none)
+
+          define_enum_none_question_method(enum_type, node)
         end
 
         unless enum_type.types["All"]?
@@ -567,6 +569,12 @@ class Crystal::TopLevelVisitor < Crystal::SemanticVisitor
     enum_type.add_def a_def
   end
 
+  def define_enum_none_question_method(enum_type, node)
+    body = Call.new(Call.new(nil, "value").at(node), "==", NumberLiteral.new(0)).at(node)
+    a_def = Def.new("none?", body: body).at(node)
+    enum_type.add_def a_def
+  end
+
   def visit(node : Assign)
     type_assign(node.target, node.value, node)
     false