From e8d50ad204cf70e424141353fe0fa9d399e40d21 Mon Sep 17 00:00:00 2001
From: Thomas Reynolds <me@tdreyno.com>
Date: Tue, 8 Nov 2011 00:00:33 -0800
Subject: [PATCH] sitemap_tree helper

---
 lib/middleman.rb                       |  3 +++
 lib/middleman/base.rb                  |  3 ++-
 lib/middleman/features/sitemap_tree.rb | 35 +++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/lib/middleman.rb b/lib/middleman.rb
index 9f12a276..06fea292 100755
--- a/lib/middleman.rb
+++ b/lib/middleman.rb
@@ -150,6 +150,9 @@ module Middleman
     
     # Automatically convert filename.html files into filename/index.html
     autoload :DirectoryIndexes,    "middleman/features/directory_indexes"
+    
+    # Organize the sitemap as a tree
+    autoload :SitemapTree,         "middleman/features/sitemap_tree"
   end
   
   EXTENSION_FILE = File.join("lib", "middleman_init.rb")
diff --git a/lib/middleman/base.rb b/lib/middleman/base.rb
index 9bcac13c..6714797f 100644
--- a/lib/middleman/base.rb
+++ b/lib/middleman/base.rb
@@ -78,7 +78,8 @@ module Middleman::Base
       app.register Middleman::CoreExtensions::FrontMatter
 
       app.set :default_features, [
-        :lorem
+        :lorem,
+        :sitemap_tree
       ]
 
       # Default layout name
diff --git a/lib/middleman/features/sitemap_tree.rb b/lib/middleman/features/sitemap_tree.rb
index 8d6b99f9..3decb117 100644
--- a/lib/middleman/features/sitemap_tree.rb
+++ b/lib/middleman/features/sitemap_tree.rb
@@ -1 +1,34 @@
-# Represent the HTML files in the site as a tree
\ No newline at end of file
+module Middleman::Features::SitemapTree
+  class << self
+    def registered(app)
+      app.helpers Helpers
+    end
+    alias :included :registered
+  end
+  
+  module Helpers
+    def sitemap_tree(regex=nil)
+      @sitemap_tree_cache = {}
+      
+      key = regex.nil? "all" : regex
+      
+      if !@sitemap_tree_cache.has_key?(key)
+        auto_hash = Hash.new{ |h,k| h[k] = Hash.new &h.default_proc }
+
+        app.sitemap.all_paths.each do |path|
+          next if !regex.nil? && !path.match(regex)
+          sub = auto_hash
+          path.split( "/" ).each{ |dir| sub[dir]; sub = sub[dir] }
+        end
+      
+        @sitemap_tree_cache[key] = auto_hash
+      end
+      
+      @sitemap_tree_cache[key]
+    end
+    
+    def html_sitemap
+      sitemap_tree(/\.html$/)
+    end
+  end
+end