From b538ce73616d3a1e91aee291fca3fafc352bfe95 Mon Sep 17 00:00:00 2001
From: Rylan Polster <rslpolster@gmail.com>
Date: Thu, 23 Jun 2022 17:18:58 -0400
Subject: [PATCH] Add `OnSystem` module to generate `on_*` methods for formulae
 and casks

---
 Library/Homebrew/cask/dsl.rb              |  20 ++++-
 Library/Homebrew/extend/on_os.rb          |  26 ------
 Library/Homebrew/extend/on_os.rbi         |   5 --
 Library/Homebrew/extend/on_system.rb      | 105 ++++++++++++++++++++++
 Library/Homebrew/extend/os/linux/on_os.rb |  10 ---
 Library/Homebrew/extend/os/mac/on_os.rb   |  10 ---
 Library/Homebrew/extend/os/on_os.rb       |   8 --
 Library/Homebrew/formula.rb               |   8 +-
 Library/Homebrew/formula.rbi              |   3 +
 Library/Homebrew/override.rb              |  53 +++++++++++
 Library/Homebrew/resource.rb              |   5 +-
 Library/Homebrew/software_spec.rb         |   7 +-
 12 files changed, 191 insertions(+), 69 deletions(-)
 delete mode 100644 Library/Homebrew/extend/on_os.rb
 delete mode 100644 Library/Homebrew/extend/on_os.rbi
 create mode 100644 Library/Homebrew/extend/on_system.rb
 delete mode 100644 Library/Homebrew/extend/os/linux/on_os.rb
 delete mode 100644 Library/Homebrew/extend/os/mac/on_os.rb
 delete mode 100644 Library/Homebrew/extend/os/on_os.rb
 create mode 100644 Library/Homebrew/override.rb

diff --git a/Library/Homebrew/cask/dsl.rb b/Library/Homebrew/cask/dsl.rb
index e4d1ac47a2..9a631484d1 100644
--- a/Library/Homebrew/cask/dsl.rb
+++ b/Library/Homebrew/cask/dsl.rb
@@ -4,6 +4,7 @@
 require "locale"
 require "lazy_object"
 require "livecheck"
+require "override"
 
 require "cask/artifact"
 
@@ -25,6 +26,8 @@ require "cask/dsl/version"
 require "cask/url"
 require "cask/utils"
 
+require "extend/on_system"
+
 module Cask
   # Class representing the domain-specific language used for casks.
   #
@@ -89,8 +92,12 @@ module Cask
       *ARTIFACT_BLOCK_CLASSES.flat_map { |klass| [klass.dsl_key, klass.uninstall_dsl_key] },
     ]).freeze
 
+    extend Predicable
+
     attr_reader :cask, :token
 
+    attr_predicate :on_system_blocks_exist?
+
     def initialize(cask)
       @cask = cask
       @token = cask.token
@@ -112,10 +119,17 @@ module Cask
     def set_unique_stanza(stanza, should_return)
       return instance_variable_get("@#{stanza}") if should_return
 
-      if !@cask.allow_reassignment && instance_variable_defined?("@#{stanza}")
-        raise CaskInvalidError.new(cask, "'#{stanza}' stanza may only appear once.")
+      unless @cask.allow_reassignment
+        if instance_variable_defined?("@#{stanza}") && !@called_in_on_system_block
+          raise CaskInvalidError.new(cask, "'#{stanza}' stanza may only appear once.")
+        end
+
+        if instance_variable_defined?("@#{stanza}_set_in_block") && @called_in_on_system_block
+          raise CaskInvalidError.new(cask, "'#{stanza}' stanza may only be overridden once.")
+        end
       end
 
+      instance_variable_set("@#{stanza}_set_in_block", true) if @called_in_on_system_block
       instance_variable_set("@#{stanza}", yield)
     rescue CaskInvalidError
       raise
@@ -306,6 +320,8 @@ module Cask
       @livecheckable == true
     end
 
+    OnSystem.setup_methods! onto: self, include_linux: false
+
     ORDINARY_ARTIFACT_CLASSES.each do |klass|
       define_method(klass.dsl_key) do |*args|
         if [*artifacts.map(&:class), klass].include?(Artifact::StageOnly) &&
diff --git a/Library/Homebrew/extend/on_os.rb b/Library/Homebrew/extend/on_os.rb
deleted file mode 100644
index aacadd608e..0000000000
--- a/Library/Homebrew/extend/on_os.rb
+++ /dev/null
@@ -1,26 +0,0 @@
-# typed: strict
-# frozen_string_literal: true
-
-module OnOS
-  extend T::Sig
-
-  # Block only executed on macOS. No-op on Linux.
-  # <pre>on_macos do
-  # # Do something Mac-specific
-  # end</pre>
-  sig { params(block: T.proc.void).void }
-  def on_macos(&block)
-    raise "No block content defined for 'on_macos' block" unless T.unsafe(block)
-  end
-
-  # Block only executed on Linux. No-op on macOS.
-  # <pre>on_linux do
-  # # Do something Linux-specific
-  # end</pre>
-  sig { params(block: T.proc.void).void }
-  def on_linux(&block)
-    raise "No block content defined for 'on_linux' block" unless T.unsafe(block)
-  end
-end
-
-require "extend/os/on_os"
diff --git a/Library/Homebrew/extend/on_os.rbi b/Library/Homebrew/extend/on_os.rbi
deleted file mode 100644
index a53474fd86..0000000000
--- a/Library/Homebrew/extend/on_os.rbi
+++ /dev/null
@@ -1,5 +0,0 @@
-# typed: strict
-
-module OnOS
-  include Kernel
-end
diff --git a/Library/Homebrew/extend/on_system.rb b/Library/Homebrew/extend/on_system.rb
new file mode 100644
index 0000000000..711bb4b187
--- /dev/null
+++ b/Library/Homebrew/extend/on_system.rb
@@ -0,0 +1,105 @@
+# typed: false
+# frozen_string_literal: true
+
+require "override"
+
+module OnSystem
+  extend T::Sig
+
+  ARCH_OPTIONS = [:intel, :arm].freeze
+  BASE_OS_OPTIONS = [:macos, :linux].freeze
+
+  module_function
+
+  sig { params(arch: Symbol).returns(T::Boolean) }
+  def arch_condition_met?(arch)
+    raise ArgumentError, "Invalid arch condition: #{arch.inspect}" unless ARCH_OPTIONS.include?(arch)
+
+    current_arch = Homebrew::Override.arch || Hardware::CPU.type
+    arch == current_arch
+  end
+
+  sig { params(os_name: Symbol, or_condition: T.nilable(Symbol)).returns(T::Boolean) }
+  def os_condition_met?(os_name, or_condition = nil)
+    if Homebrew::EnvConfig.simulate_macos_on_linux?
+      return false if os_name == :linux
+      return true if [:macos, *MacOS::Version::SYMBOLS.keys].include?(os_name)
+    end
+
+    if BASE_OS_OPTIONS.include?(os_name)
+      if Homebrew::Override.none?
+        return OS.linux? if os_name == :linux
+        return OS.mac? if os_name == :macos
+      end
+
+      return Homebrew::Override.send("#{os_name}?")
+    end
+
+    raise ArgumentError, "Invalid OS condition: #{os_name.inspect}" unless MacOS::Version::SYMBOLS.key?(os_name)
+
+    if or_condition.present? && [:or_newer, :or_older].exclude?(or_condition)
+      raise ArgumentError, "Invalid OS `or_*` condition: #{or_condition.inspect}"
+    end
+
+    base_os = MacOS::Version.from_symbol(os_name)
+    current_os = MacOS::Version.from_symbol(Homebrew::Override.os || MacOS.version.to_sym)
+
+    return current_os >= base_os if or_condition == :or_newer
+    return current_os <= base_os if or_condition == :or_older
+
+    current_os == base_os
+  end
+
+  sig { params(method_name: Symbol).returns(Symbol) }
+  def condition_from_method_name(method_name)
+    method_name.to_s.sub(/^on_/, "").to_sym
+  end
+
+  sig { params(onto: Class, include_linux: T::Boolean).void }
+  def setup_methods!(onto:, include_linux: true)
+    ARCH_OPTIONS.each do |arch|
+      onto.define_method("on_#{arch}") do |&block|
+        @on_system_blocks_exist = true
+
+        return unless OnSystem.arch_condition_met? OnSystem.condition_from_method_name(__method__)
+
+        @called_in_on_system_block = true
+        result = block.call
+        @called_in_on_system_block = false
+
+        result
+      end
+    end
+
+    if include_linux
+      BASE_OS_OPTIONS.each do |base_os|
+        onto.define_method("on_#{base_os}") do |&block|
+          @on_system_blocks_exist = true
+
+          return unless OnSystem.os_condition_met? OnSystem.condition_from_method_name(__method__)
+
+          @called_in_on_system_block = true
+          result = block.call
+          @called_in_on_system_block = false
+
+          result
+        end
+      end
+    end
+
+    MacOS::Version::SYMBOLS.each_key do |os_name|
+      onto.define_method("on_#{os_name}") do |or_condition = nil, &block|
+        @on_system_blocks_exist = true
+
+        os_condition = OnSystem.condition_from_method_name __method__
+        return unless OnSystem.os_condition_met? os_condition, or_condition
+
+        @called_in_on_system_block = true
+        result = block.call
+        @called_in_on_system_block = false
+
+        result
+      end
+    end
+  end
+end
diff --git a/Library/Homebrew/extend/os/linux/on_os.rb b/Library/Homebrew/extend/os/linux/on_os.rb
deleted file mode 100644
index 4373a45ab2..0000000000
--- a/Library/Homebrew/extend/os/linux/on_os.rb
+++ /dev/null
@@ -1,10 +0,0 @@
-# typed: strict
-# frozen_string_literal: true
-
-module OnOS
-  def on_linux(&block)
-    raise "No block content defined for 'on_linux' block" unless T.unsafe(block)
-
-    yield
-  end
-end
diff --git a/Library/Homebrew/extend/os/mac/on_os.rb b/Library/Homebrew/extend/os/mac/on_os.rb
deleted file mode 100644
index 78759c1a9d..0000000000
--- a/Library/Homebrew/extend/os/mac/on_os.rb
+++ /dev/null
@@ -1,10 +0,0 @@
-# typed: strict
-# frozen_string_literal: true
-
-module OnOS
-  def on_macos(&block)
-    raise "No block content defined for 'on_macos' block" unless T.unsafe(block)
-
-    yield
-  end
-end
diff --git a/Library/Homebrew/extend/os/on_os.rb b/Library/Homebrew/extend/os/on_os.rb
deleted file mode 100644
index c74ea9d459..0000000000
--- a/Library/Homebrew/extend/os/on_os.rb
+++ /dev/null
@@ -1,8 +0,0 @@
-# typed: strict
-# frozen_string_literal: true
-
-if OS.mac? || Homebrew::EnvConfig.simulate_macos_on_linux?
-  require "extend/os/mac/on_os"
-elsif OS.linux?
-  require "extend/os/linux/on_os"
-end
diff --git a/Library/Homebrew/formula.rb b/Library/Homebrew/formula.rb
index 25ff2c3920..250583b465 100644
--- a/Library/Homebrew/formula.rb
+++ b/Library/Homebrew/formula.rb
@@ -28,7 +28,7 @@ require "tab"
 require "mktemp"
 require "find"
 require "utils/spdx"
-require "extend/on_os"
+require "extend/on_system"
 require "api"
 
 # A formula provides instructions and metadata for Homebrew to install a piece
@@ -64,11 +64,12 @@ class Formula
   include Utils::Shebang
   include Utils::Shell
   include Context
-  include OnOS
   extend Forwardable
   extend Cachable
   extend Predicable
 
+  OnSystem.setup_methods! onto: self
+
   # @!method inreplace(paths, before = nil, after = nil)
   # @see Utils::Inreplace.inreplace
 
@@ -2470,7 +2471,6 @@ class Formula
   # The methods below define the formula DSL.
   class << self
     include BuildEnvironment::DSL
-    include OnOS
 
     def method_added(method)
       super
@@ -2483,6 +2483,8 @@ class Formula
       end
     end
 
+    OnSystem.setup_methods! onto: self
+
     # The reason for why this software is not linked (by default) to
     # {::HOMEBREW_PREFIX}.
     # @private
diff --git a/Library/Homebrew/formula.rbi b/Library/Homebrew/formula.rbi
index c573f64b3c..a47b034dee 100644
--- a/Library/Homebrew/formula.rbi
+++ b/Library/Homebrew/formula.rbi
@@ -48,4 +48,7 @@ class Formula
 
   def env; end
   def conflicts; end
+
+  # This method is generated by `OnSystem::setup_methods!`
+  def self.on_macos(&block); end
 end
diff --git a/Library/Homebrew/override.rb b/Library/Homebrew/override.rb
new file mode 100644
index 0000000000..7b6765ac5d
--- /dev/null
+++ b/Library/Homebrew/override.rb
@@ -0,0 +1,53 @@
+# typed: true
+# frozen_string_literal: true
+
+module Homebrew
+  # Helper module for overriding machine-specific methods.
+  #
+  # @api private
+  class Override
+    class << self
+      extend T::Sig
+
+      ARCH_OPTIONS = [:intel, :arm64].freeze
+      OS_OPTIONS = [:macos, :linux].freeze
+
+      attr_reader :os, :arch
+
+      sig { params(new_os: Symbol).void }
+      def os=(new_os)
+        os_options = [:macos, :linux, *MacOS::Version::SYMBOLS.keys]
+        raise "Unknown OS: #{new_os}" unless os_options.include?(new_os)
+
+        @os = new_os
+      end
+
+      sig { params(new_arch: Symbol).void }
+      def arch=(new_arch)
+        raise "New arch must be :arm or :intel" unless [:arm, :intel].include?(new_arch)
+
+        @arch = new_arch
+      end
+
+      sig { void }
+      def clear
+        @os = @arch = nil
+      end
+
+      sig { returns(T::Boolean) }
+      def none?
+        @os.nil? && @arch.nil?
+      end
+
+      sig { returns(T::Boolean) }
+      def macos?
+        [:macos, *MacOS::Version::SYMBOLS.keys].include?(@os)
+      end
+
+      sig { returns(T::Boolean) }
+      def linux?
+        @os == :linux
+      end
+    end
+  end
+end
diff --git a/Library/Homebrew/resource.rb b/Library/Homebrew/resource.rb
index 49d9e362a8..9cbd64bc7f 100644
--- a/Library/Homebrew/resource.rb
+++ b/Library/Homebrew/resource.rb
@@ -5,7 +5,7 @@ require "download_strategy"
 require "checksum"
 require "version"
 require "mktemp"
-require "extend/on_os"
+require "extend/on_system"
 
 # Resource is the fundamental representation of an external resource. The
 # primary formula download, along with other declared resources, are instances
@@ -17,7 +17,8 @@ class Resource
 
   include Context
   include FileUtils
-  include OnOS
+
+  OnSystem.setup_methods! onto: self
 
   attr_reader :mirrors, :specs, :using, :source_modified_time, :patches, :owner
   attr_writer :version
diff --git a/Library/Homebrew/software_spec.rb b/Library/Homebrew/software_spec.rb
index 437af9f0a8..61a0b8d463 100644
--- a/Library/Homebrew/software_spec.rb
+++ b/Library/Homebrew/software_spec.rb
@@ -12,19 +12,20 @@ require "utils/bottles"
 require "patch"
 require "compilers"
 require "os/mac/version"
-require "extend/on_os"
+require "extend/on_system"
 
 class SoftwareSpec
   extend T::Sig
 
   extend Forwardable
-  include OnOS
 
   PREDEFINED_OPTIONS = {
     universal: Option.new("universal", "Build a universal binary"),
     cxx11:     Option.new("c++11",     "Build using C++11 mode"),
   }.freeze
 
+  OnSystem.setup_methods! onto: self
+
   attr_reader :name, :full_name, :owner, :build, :resources, :patches, :options, :deprecated_flags,
               :deprecated_options, :dependency_collector, :bottle_specification, :compiler_failures,
               :uses_from_macos_elements
@@ -583,7 +584,7 @@ class BottleSpecification
 end
 
 class PourBottleCheck
-  include OnOS
+  OnSystem.setup_methods! onto: self
 
   def initialize(formula)
     @formula = formula