11 files changed, 112 insertions, 13 deletions
diff --git a/lib/yard/mruby/code_objects.rb b/lib/yard/mruby/code_objects.rb
index 07dbf8b..bd41cef 100644
--- a/lib/yard/mruby/code_objects.rb
+++ b/lib/yard/mruby/code_objects.rb
@@ -1,4 +1,5 @@
 require_relative 'code_objects/headers_root'
+require_relative 'code_objects/header_base_object'
 require_relative 'code_objects/header_object'
 require_relative 'code_objects/function_object'
 require_relative 'code_objects/define_object'
diff --git a/lib/yard/mruby/code_objects/define_object.rb b/lib/yard/mruby/code_objects/define_object.rb
index 2c05d91..4220acd 100644
--- a/lib/yard/mruby/code_objects/define_object.rb
+++ b/lib/yard/mruby/code_objects/define_object.rb
@@ -1,7 +1,7 @@
 module YARD::MRuby::CodeObjects
 
   # A DefineObject represents a MRuby C API define macro declaration inside a header inside an include directory
-  class DefineObject < YARD::CodeObjects::Base
+  class DefineObject < HeaderBaseObject
   end
 end
 
diff --git a/lib/yard/mruby/code_objects/function_object.rb b/lib/yard/mruby/code_objects/function_object.rb
index 7d85439..ecf7dd5 100644
--- a/lib/yard/mruby/code_objects/function_object.rb
+++ b/lib/yard/mruby/code_objects/function_object.rb
@@ -1,7 +1,7 @@
 module YARD::MRuby::CodeObjects
 
   # A FunctionObject represents a MRuby C API function declaration inside a header inside an include directory
-  class FunctionObject < YARD::CodeObjects::Base
+  class FunctionObject < HeaderBaseObject
 
     # Returns the list of parameters parsed out of the method signature
     # with their default values.
@@ -14,15 +14,6 @@ module YARD::MRuby::CodeObjects
       super
     end
 
-    def header
-      self.namespace
-    end
-
-    # Function's shouln't be namespaced
-    def path
-      self.name.to_s
-    end
-
     def attr_info
       nil
     end
diff --git a/lib/yard/mruby/code_objects/header_base_object.rb b/lib/yard/mruby/code_objects/header_base_object.rb
new file mode 100644
index 0000000..403db3d
--- /dev/null
+++ b/lib/yard/mruby/code_objects/header_base_object.rb
@@ -0,0 +1,15 @@
+module YARD::MRuby::CodeObjects
+
+  # A FunctionObject represents a MRuby C API function declaration inside a header inside an include directory
+  class HeaderBaseObject < YARD::CodeObjects::Base
+
+    def header
+      self.namespace
+    end
+
+    # Header objects's shouln't be namespaced
+    def path
+      self.name.to_s
+    end
+  end
+end
diff --git a/lib/yard/mruby/handlers/c/header.rb b/lib/yard/mruby/handlers/c/header.rb
index 60848cd..9c955f2 100644
--- a/lib/yard/mruby/handlers/c/header.rb
+++ b/lib/yard/mruby/handlers/c/header.rb
@@ -1,3 +1,4 @@
 require_relative 'header/base'
 require_relative 'header/function_handler'
+require_relative 'header/define_handler'
 require_relative 'header/header_decl_handler'
diff --git a/lib/yard/mruby/handlers/c/header/define_handler.rb b/lib/yard/mruby/handlers/c/header/define_handler.rb
index e69de29..1a061db 100644
--- a/lib/yard/mruby/handlers/c/header/define_handler.rb
+++ b/lib/yard/mruby/handlers/c/header/define_handler.rb
@@ -0,0 +1,23 @@
+module YARD::MRuby::Handlers::C::Header
+  class DefineHandler < Base
+    MATCH = /
+      \#\s*define\s*(\w+)
+    /mx
+
+    handles MATCH
+    statement_class YARD::MRuby::Parser::C::DirectiveStatement
+
+    process do
+      header = self.header(statement.file)
+
+      statement.declaration.scan(MATCH) do |match|
+        register DefineObject.new(header, match.first) do |obj|
+          if statement.comments
+            register_docstring(obj, statement.comments.source, statement)
+          end
+        end
+      end
+
+    end
+  end
+end
diff --git a/lib/yard/mruby/parser/c.rb b/lib/yard/mruby/parser/c.rb
index 4d95fa1..13f5ef7 100644
--- a/lib/yard/mruby/parser/c.rb
+++ b/lib/yard/mruby/parser/c.rb
@@ -1,3 +1,4 @@
 require_relative 'c/parser'
 require_relative 'c/source_parser'
 require_relative 'c/header_parser'
+require_relative 'c/directive_statement'
diff --git a/lib/yard/mruby/parser/c/directive_statement.rb b/lib/yard/mruby/parser/c/directive_statement.rb
new file mode 100644
index 0000000..b146d17
--- /dev/null
+++ b/lib/yard/mruby/parser/c/directive_statement.rb
@@ -0,0 +1,6 @@
+module YARD::MRuby::Parser::C
+  class DirectiveStatement < YARD::Parser::C::Statement
+    attr_accessor :declaration
+    attr_accessor :comments
+  end
+end
diff --git a/lib/yard/mruby/parser/c/header_parser.rb b/lib/yard/mruby/parser/c/header_parser.rb
index 598fe21..2b2dea7 100644
--- a/lib/yard/mruby/parser/c/header_parser.rb
+++ b/lib/yard/mruby/parser/c/header_parser.rb
@@ -1,5 +1,46 @@
 module YARD::MRuby::Parser::C
   class HeaderParser < Parser
+
+    # Consumes a directive and generates a DirectiveStatement
+    def consume_directive
+      super if @in_body_statements
+
+      @newline = false
+      start = @index
+      line = @line
+      statement = DirectiveStatement.new(nil, @file, line)
+      @statements << statement
+      attach_comment(statement)
+
+      multiline = false
+      advance_loop do
+        chr = char
+        case chr
+        when '\\'; multiline=true; advance
+        when /\s/; consume_whitespace
+        else advance
+        end
+
+        if @newline
+          if multiline
+            multiline = false
+          else
+            break
+          end
+        end
+      end
+
+      decl = @content[start...@index]
+
+      statement.declaration = decl
+    end
+
+    def consume_body_statements
+      @in_body_statements = true
+      result = super
+      @in_body_statements = false
+      result
+    end
   end
 
   #
diff --git a/spec/handlers/c/header/define_handler_spec.rb b/spec/handlers/c/header/define_handler_spec.rb
new file mode 100644
index 0000000..c37e347
--- /dev/null
+++ b/spec/handlers/c/header/define_handler_spec.rb
@@ -0,0 +1,22 @@
+require_relative 'spec_helper'
+
+describe YARD::MRuby::Handlers::C::Header::DefineHandler do
+  it "should register defines" do
+    header_line <<-eof
+      # define MRB_FOO bar
+    eof
+
+    puts Registry.send(:thread_local_store).inspect
+    expect(Registry.at('MRB_FOO')).not_to be_nil
+  end
+
+  it "should find docstrings attached to defines" do
+    header_line <<-eof
+      /* DOCSTRING */
+      # define MRB_FOO bar
+    eof
+
+    define = Registry.at('MRB_FOO')
+    expect(define.docstring).to eq 'DOCSTRING'
+  end
+end
diff --git a/spec/handlers/c/header/function_handler_spec.rb b/spec/handlers/c/header/function_handler_spec.rb
index 2270584..dc42c91 100644
--- a/spec/handlers/c/header/function_handler_spec.rb
+++ b/spec/handlers/c/header/function_handler_spec.rb
@@ -6,8 +6,6 @@ describe YARD::MRuby::Handlers::C::Header::FunctionHandler do
       MRB_API void mrb_foo( void );
     eof
     expect(Registry.at('mrb_foo')).not_to be_nil
-
-    # puts Registry.send(:thread_local_store).inspect
   end
 
   it "should find docstrings attached to functions" do