Script for calculating probabilities for custom dice system

dice_system.rb

require 'set'

class Or
  def initialize *ands
    @ands = ands
  end

  def to_a
    @ands
  end

  def inspect
    fs self
  end
end

# get a set of faces as a string
def ds d
  "{#{d.map { |f| fs f }.join(', ')}}"
end

# get a set of concrete faces as a string
def dcs d
  "{#{d.map { |f| cfs f }.join(', ')}}"
end

# get a face as a string
def fs f
  if f.is_a? Or
    f.to_a.map { |a| fs a }.join(?/)
  else
    f.map { |cf| cfs cf }.join(?+)
  end
end

# get a concrete face as a string
def cfs cf
  "#{cf[0] == :wild ? ?? : cf[0].to_s.upcase}#{cf[1]}"
end

# for a set of faces, generate the array of concrete faces one could choose (no And/Or/blank)
def literals d
  ors, ands = d.partition { |f| f.is_a? Or }

  # no choices to make
  if ors.empty?
    # flatten since all faces now count
    return [ands.flatten(1)]
  end

  ors = ors.map { |o| o.to_a }

  # all possible Or choices
  ors[0].product(*ors[1..-1]).map do |o|
    # combined with the ands, flattened since all faces now count
    (o + ands).flatten(1).sort
  end.uniq
end

# for a set of literal faces (no And/Or/blank), get an array of concrete test results one could choose by combining faces
# returns an array of concrete skills
def concrete_results literal

  skills = literal.group_by(&:first).map do |s, fs|
    counts = fs.map(&:last)

    # all the different ways you could combine your faces for this skill
    opts = multiset_partitions(counts).map do |partition|
      partition.map(&:sum).sort
    end.uniq

    opts.map { |opt|
      opt.map { |o| [s, o] }
    }
  end

  # all the different ways you could combine your choices for all skills
  skills[0].product(*skills[1..-1]).map { |choice|
    # flatten because we've decided how each skill is/isn't combined, so this is a final result
    choice.flatten(1)
  }
end

def parse_face str
  if str.strip =~ /\A([PSTCI?])(\d+)\Z/
    [$1 == ?? ? :wild : $1.to_sym, Integer($2, 10)]
  else
    abort "Unrecognized die face format: #{str.strip}"
  end
end

$skills = [:P, :S, :T, :C, :I]

# parse a string into a die
#
# a die is as an arry of Ors
#
# an Or contains zero or more Ands, representing / (you can choose to use one of the Ands)
#
# an And is an array of concrete faces, representing a + (all concrete faces can be used)
#
# a concrete face is [skill, value] where skill is a symbol (including :wild as a skill)
#
# even though internally we support nesting And inside of Or, the parser does not
#
# this transforms wilds into a normalized form where Ors are used to represent the specific choices you might make for the wild
def parse str
  # extra space since ',' should be parsed as two blanks
  str.gsub(?,, ', ').split(?,).map do |face|
    case face.strip
    when ''
      []
    when /\//
      Or.new(*face.split(?\/).map do |or_s|
        f = parse_face(or_s)
        if f[0] == :wild
          # note the extra wrapping since all of these get flattened into this whole list
          $skills.map { |s| [[s, f[1]]] } + [[f]]
        else
          [[f]]
        end
      end.flatten(1))
    when /\+/
      faces = face.split(?+).map { |s| parse_face(s) }
      wilds, non = faces.partition { |f| f[0] == :wild }

      if wilds.empty?
        next faces
      end

      # explode all the wilds
      wilds = wilds.map { |w| $skills.map { |s| [s, w[1]] } + [w] }

      # each wild can be chosen independently
      Or.new(*wilds[0].product(*wilds[1..-1]).map do |ws|
        ws + non
      end)
    else
      f = parse_face face
      if f[0] == :wild
        Or.new(*($skills.map { |s| [[s, f[1]]] } + [[f]]))
      else
        [f]
      end
    end
  end
end

# given a multiset represented as an array in any order e.g., [1, 2, 4, 4, 2, 2]
# generate all partitions:
#
# [
#   [[1, 2, 2, 2, 4, 4]],
#   [[1, 2, 2, 2, 4], [4]],
#   ...
#   [[1], [2], [2], [2], [4, 4]],
#   [[1], [2], [2], [2], [4], [4]]
# ]
#
# from The Art of Computer Programming Volume 4A Section 7.2.1.5 pg 429 Algorithm M
def multiset_partitions multiset
  counts = {}
  multiset.each do |x|
    counts[x] ||= 0
    counts[x] += 1
  end

  components = counts.keys.sort
  m = components.length # number of distinct values in multiset
  n = [] # n[i] = number of occurences of components[i] in multiset
  c = [] # component number: there are v[i] occurences of components[c[i]] in this part of the partition
  u = [] # there are u[i] occurences of components[c[i]] yet unpartitioned (counting v[i] as unpartitioned)
  v = [] # see c
  l = 0 # current partition has l+1 parts
  a = l # start of current stack frame
  b = m # start of next stack frame
  f = [a, b] # array of starts of stack frames

  partitions = [] # records output as we go

  # M1 init
  (0...m).each do |j|
    n[j] = counts[components[j]]
    c[j] = j
    u[j] = v[j] = n[j]
  end

  # M2 subtract v from u
  # at this point we want to find all partitions of the vector u in the current frame, into parts that are lexicographically <= v
  # first we use v itself
  loop do
    loop do
      j = a
      k = b
      x = false # has v changed?
      while j < b
        u[k] = u[j] - v[j]
        if u[k] == 0
          x = true
          j += 1
        elsif !x
          c[k] = c[j]
          v[k] = [v[j], u[k]].min
          x = u[k] < v[j]
          j += 1
          k += 1
        else
          c[k] = c[j]
          v[k] = u[k]
          j += 1
          k += 1
        end
      end

      # M3 push if nonzero
      if k > b
        a = b
        b = k
        l += 1
        f[l + 1] = b
        # don't break. return to M2
      else
        break # proceed to M4
      end
    end

    # M4 visit a partition
    partitions << (0..l).map do |k|
      part = []
      (f[k]...f[k + 1]).each do |j|
        v[j].times { part << components[c[j]] }
      end
      part
    end

    # M5 decrease v
    loop do
      j = b - 1
      while v[j] == 0
        j -= 1
      end

      if j == a && v[j] == 1
        # M6 backtrack
        return partitions if l == 0
        l = l - 1
        b = a
        a = f[l]
        # don't break. return to M5
      else
        v[j] -= 1
        ((j + 1)...b).each do |k|
          v[k] = u[k]
        end
        break # return to M2
      end
    end
  end
end


# given any number of dice, show the chance of getting each possible result
def roll test, *dice
  total = dice.map(&:length).reduce(:*)
  success = 0

  results = {}

  dice[0].product(*dice[1..-1]) do |roll_faces|
    passes = false

    lits = literals(roll_faces)

    res = []

    lits.each do |lit|
      res.concat concrete_results(lit)
    end

    res.uniq!

    res.each do |result|
      if !passes && test && passes?(result, test)
        passes = true
      end
      results[result] ||= 0
      results[result] += 1
    end

    if passes
      success += 1
    end
  end


  if test
    puts "Possible results that match marked with *"
  else
    puts "Showing all results. Note that probabilities do not add up!"
  end

  results.to_a.sort_by(&:last).reverse.each do |result, count|
    passes = ' '
    if test && passes?(result, test)
      passes = ?*
    end
    puts "%3d%%\t#{passes}#{dcs(result)}" % (100.0 * count / total)
  end

  if test
    puts "%3d%% chance of success" % (100.0 * success / total)
  else
    puts "No test specified. Use -t to see a specific chance of success"
  end
end

# parse a test string into a skill test
# produces an array of [skill, value] pairs
def parse_test str
  str = str.strip
  if str =~ /\A([PSTCI?A](\d+))+\Z/
    test = []
    i = 0
    loop do
      skill = str[i] == ?? ? :wild : str[i].to_sym
      start = i + 1
      i += 1
      while i < str.length && str[i] =~ /\d/
        i += 1
      end
      test << [skill, Integer(str[start...i], 10)]
      break if i >= str.length
    end
    return test
  else
    abort "unrecognized test format: #{str}"
  end
end

# check if a concrete result passes a test
#
# note that a concrete result has already decided if ? is being used as another skill, so ? _only matches_ ?
def passes? result, test
  result = result.group_by(&:first).map do |s, fs|
    [s, fs.map(&:last).sort]
  end.to_h
  test = test.group_by(&:first)

  ($skills + [:wild]).each do |s|
    next unless test[s]
    return false unless result[s]

    required = test[s].map(&:last).sort

    until required.empty?
      need = required.shift
      at = result[s].index { |c| c >= need }
      return false unless at
      result[s].delete_at at
    end
  end

  # if we matched all skills+wilds and there are no Anys, we're done
  return true unless test[:A]

  required = test[:A].map(&:last).sort
  # we don't care about skills anymore, we just need counts
  have = result.values.reduce([], :concat).sort

  until required.empty?
    need = required.shift
    at = have.index { |c| c >= need }
    return false unless at
    have.delete_at at
  end

  return true
end

dice_arg_end = -1

if test_index = ARGV.index('-t')
  dice_arg_end = test_index - 1
end

if test_index && test_index < 2
  abort "you must specify all dice with -d before -t"
end

dice = ARGV[0..dice_arg_end].join(' ').split(/\s*-d/)

unless dice[0]
  abort "usage: specify each die with -d. follow with -t if you want to check a specific test chance"
end

if dice[0] != ''
  abort "unrecognized argument #{dice[0]}. you must specify all dice with -d"
end

puts "Normalizing dice:"
dice = dice[1..-1].map { |d|
  x = parse(d)
  puts ds(x)
  x
}

test = nil
if test_index
  test = parse_test(ARGV[(test_index + 1)..-1].join.split(' ').join)
end

roll test, *dice