From 7dae24894df275549c37b2d810541a7e7ea306b3 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Wed, 24 Jul 2019 16:01:20 -0400
Subject: Implement `freeze` option for Pysch.load

---
 lib/psych/visitors/to_ruby.rb | 18 +++++++++++-------
 test/psych/test_psych.rb      | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index 68452e5..a922f90 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -32,16 +32,18 @@ module Psych
 
       def accept target
         result = super
-        return result if @domain_types.empty? || !target.tag
 
-        key = target.tag.sub(/^[!\/]*/, '').sub(/(,\d+)\//, '\1:')
-        key = "tag:#{key}" unless key =~ /^(?:tag:|x-private)/
+        unless @domain_types.empty? || !target.tag
+          key = target.tag.sub(/^[!\/]*/, '').sub(/(,\d+)\//, '\1:')
+          key = "tag:#{key}" unless key =~ /^(?:tag:|x-private)/
 
-        if @domain_types.key? key
-          value, block = @domain_types[key]
-          return block.call value, result
+          if @domain_types.key? key
+            value, block = @domain_types[key]
+            result = block.call value, result
+          end
         end
 
+        result = deduplicate(result).freeze if @freeze
         result
       end
 
@@ -341,7 +343,7 @@ module Psych
           key = accept(k)
           if @symbolize_names
             key = key.to_sym
-          else
+          elsif !@freeze
             key = deduplicate(key)
           end
           val = accept(v)
@@ -378,6 +380,8 @@ module Psych
       if RUBY_VERSION < '2.7'
         def deduplicate key
           if key.is_a?(String)
+            # It is important to untaint the string, otherwise it won't
+            # be deduplicated into an fstring, but simply frozen.
             -(key.untaint)
           else
             key
diff --git a/test/psych/test_psych.rb b/test/psych/test_psych.rb
index eeadc86..8fd9d27 100644
--- a/test/psych/test_psych.rb
+++ b/test/psych/test_psych.rb
@@ -192,6 +192,22 @@ class TestPsych < Psych::TestCase
     assert_equal({ 'hello' => 'world' }, got)
   end
 
+  def test_load_freeze
+    data = Psych.load("--- {foo: ['a']}", freeze: true)
+    assert_predicate data, :frozen?
+    assert_predicate data['foo'], :frozen?
+    assert_predicate data['foo'].first, :frozen?
+  end
+
+  def test_load_freeze_deduplication
+    unless String.method_defined?(:-@) && (-("a" * 20)).equal?((-("a" * 20)))
+      skip "This Ruby implementation doesn't support string deduplication"
+    end
+
+    data = Psych.load("--- ['a']", freeze: true)
+    assert_same 'a', data.first
+  end
+
   def test_load_default_fallback
     assert_equal false, Psych.load("")
   end
-- 
cgit v1.2.3