summaryrefslogtreecommitdiff
path: root/lib/gitlab/utils/strong_memoize.rb
blob: 50b8428113dd3a52b6ef18596d1dd3ce666ac134 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# frozen_string_literal: true

module Gitlab
  module Utils
    module StrongMemoize
      # Instead of writing patterns like this:
      #
      #     def trigger_from_token
      #       return @trigger if defined?(@trigger)
      #
      #       @trigger = Ci::Trigger.find_by_token(params[:token].to_s)
      #     end
      #
      # We could write it like:
      #
      #     include Gitlab::Utils::StrongMemoize
      #
      #     def trigger_from_token
      #       strong_memoize(:trigger) do
      #         Ci::Trigger.find_by_token(params[:token].to_s)
      #       end
      #     end
      #
      # Or like:
      #
      #     include Gitlab::Utils::StrongMemoize
      #
      #     def trigger_from_token
      #       Ci::Trigger.find_by_token(params[:token].to_s)
      #     end
      #     strong_memoize_attr :trigger_from_token
      #
      #     strong_memoize_attr :enabled?, :enabled
      #     def enabled?
      #       Feature.enabled?(:some_feature)
      #     end
      #
      def strong_memoize(name)
        key = ivar(name)

        if instance_variable_defined?(key)
          instance_variable_get(key)
        else
          instance_variable_set(key, yield)
        end
      end

      def strong_memoized?(name)
        instance_variable_defined?(ivar(name))
      end

      def clear_memoization(name)
        key = ivar(name)
        remove_instance_variable(key) if instance_variable_defined?(key)
      end

      module StrongMemoizeClassMethods
        def strong_memoize_attr(method_name, member_name = nil)
          member_name ||= method_name

          if method_defined?(method_name) || private_method_defined?(method_name)
            StrongMemoize.send( # rubocop:disable GitlabSecurity/PublicSend
              :do_strong_memoize, self, method_name, member_name)
          else
            StrongMemoize.send( # rubocop:disable GitlabSecurity/PublicSend
              :queue_strong_memoize, self, method_name, member_name)
          end
        end

        def method_added(method_name)
          super

          if member_name = StrongMemoize
            .send(:strong_memoize_queue, self).delete(method_name) # rubocop:disable GitlabSecurity/PublicSend
            StrongMemoize.send( # rubocop:disable GitlabSecurity/PublicSend
              :do_strong_memoize, self, method_name, member_name)
          end
        end
      end

      def self.included(base)
        base.singleton_class.prepend(StrongMemoizeClassMethods)
      end

      private

      # Convert `"name"`/`:name` into `:@name`
      #
      # Depending on a type ensure that there's a single memory allocation
      def ivar(name)
        if name.is_a?(Symbol)
          name.to_s.prepend("@").to_sym
        elsif name.is_a?(String)
          :"@#{name}"
        else
          raise ArgumentError, "Invalid type of '#{name}'"
        end
      end

      class <<self
        private

        def strong_memoize_queue(klass)
          klass.instance_variable_get(:@strong_memoize_queue) || klass.instance_variable_set(:@strong_memoize_queue, {})
        end

        def queue_strong_memoize(klass, method_name, member_name)
          strong_memoize_queue(klass)[method_name] = member_name
        end

        def do_strong_memoize(klass, method_name, member_name)
          method = klass.instance_method(method_name)

          # Methods defined within a class method are already public by default, so we don't need to
          # explicitly make them public.
          scope = %i[private protected].find do |scope|
            klass.send("#{scope}_instance_methods") # rubocop:disable GitlabSecurity/PublicSend
              .include? method_name
          end

          klass.define_method(method_name) do |*args, &block|
            strong_memoize(member_name) do
              method.bind_call(self, *args, &block)
            end
          end

          klass.send(scope, method_name) if scope # rubocop:disable GitlabSecurity/PublicSend
        end
      end
    end
  end
end