module GraphQL::Client::CollocatedEnforcement

def allow_noncollocated_callers
  Thread.current[:query_result_caller_location_ignore] = true
  yield
ensure
  Thread.current[:query_result_caller_location_ignore] = nil
end

def enforce_collocated_callers(mod, methods, path)
  mod.prepend(Module.new do
    methods.each do |method|
      define_method(method) do |*args, &block|
        location = caller_locations(1, 1)[0]
        CollocatedEnforcement.verify_collocated_path(location, path, method) do
          super(*args, &block)
        end
      end
    end
  end)
end

def verify_collocated_path(location, path, method = "method")
  return yield if Thread.current[:query_result_caller_location_ignore]
  if (location.path != path) && !(WHITELISTED_GEM_NAMES.any? { |g| location.path.include?("gems/#{g}") })
    error = NonCollocatedCallerError.new("#{method} was called outside of '#{path}' https://github.com/github-community-projects/graphql-client/blob/master/guides/collocated-call-sites.md")
    error.set_backtrace(caller(2))
    raise error
  end
  begin
    Thread.current[:query_result_caller_location_ignore] = true
    yield
  ensure
    Thread.current[:query_result_caller_location_ignore] = nil
  end
end