model: allow using TreeNode objects as indices in reference sets

1 files changed, 8 insertions, 8 deletions

diff --git a/acf2/model/set.lua b/acf2/model/set.lua
index 4fcbc24..8d255d9 100644
--- a/acf2/model/set.lua
+++ b/acf2/model/set.lua
@@ -21,8 +21,12 @@ function M.Set:init(context, params)
    object.super(self, M.Set):init(context, params, 'set')
 
    local function find(value)
+      value = node.BoundMember(
+	 self, pth.wildcard, params.field
+      ):normalize(value)
+
       for i, member in node.pairs(self) do
-	 if member == value then return i end
+	 if member == value then return i, value end
       end
    end
 
@@ -30,9 +34,9 @@ function M.Set:init(context, params)
 
    function mt.get(k, options)
       options = setdefaults(options or {}, {dereference=true})
-      local i = find(k)
+      local i, v = find(k)
       if i then return mt.load(i, options) end
-      if options.create then return k end
+      if options.create then return v end
    end
 
    function mt.__newindex(t, k, v)
@@ -42,11 +46,7 @@ function M.Set:init(context, params)
       mt.save(i, nil)
    end
 
-   function mt.contains(v)
-      return find(
-	 node.BoundMember(self, pth.wildcard, params.field):normalize(v)
-      ) and true or false
-   end
+   function mt.contains(v) return find(v) and true or false end
 
    local insert = mt.insert
    function mt.insert(v) if not mt.contains(v) then insert(v) end end