ruby SPF/Dijkstra implementation

spf.rb

#!/usr/bin/env ruby

## Example run: (omitting DST will show cost+path to all dsts)
## [ytti@eng0.dllstx09.us.to.gin.ntt.net ~/git/misc]% ./spf.rb r23.snjsca04.us.bb r20.frnkge04.de.bb
## ---
## r20.frnkge04.de.bb:
##   cost: 155152
##   path:
##   - r23.snjsca04.us.bb
##   - r22.snjsca04.us.bb
##   - r23.lsanca07.us.bb
##   - r22.dllstx09.us.bb
##   - r20.atlnga05.us.bb
##   - r22.asbnva02.us.bb
##   - r21.frnkge03.de.bb
## [ytti@eng0.dllstx09.us.to.gin.ntt.net ~/git/misc]%

#require 'pry'

class  SPF
  attr_accessor :nodes, :edges, :tmp
  def initialize nodes=[], edges=[]
    @nodes = nodes
    @edges = edges
  end

  def add_node node
    @nodes << node
  end

  def add_edge src, dst, metric
    @edges << { src: src, dst: dst, metric: metric }
  end

  def add_edge_symmetric src, dst, metric
    add_edge src, dst, metric
    add_edge dst, src, metric
  end

  def cost src, dst
    spf(src, dst)[dst][:cost]
  end

  def path src, dst
    spf(src, dst)[dst][:path]
  end

  def ero_cost path
    return Float::INFINITY unless path
    path.each_cons(2).inject(0) do |cost, (src, dst)|
      cost + (edge(src, dst)[:metric] rescue Float::INFINITY)
    end
  end

  def spf src, dst=nil
    nodes = @nodes.dup
    prevs = {}
    costs = Hash.new Float::INFINITY
    costs[src] = 0

    while nearest_node = nodes.delete(nodes.min_by { |node| costs[node] })
      return spf_result(src, dst, costs, prevs) if nearest_node == dst
      edges(nearest_node).each do |edge|
        neighbor, metric = edge[:dst], edge[:metric]
        metric += costs[nearest_node]
        if metric < costs[neighbor]
          costs[neighbor] = metric
          prevs[neighbor] = nearest_node
        end
      end
    end
    spf_result src, dst, costs, prevs
  end

  def edges src
    @edges.select { |e| e[:src]==src }
  end

  def edge src, dst
    @edges.select { |e| e[:src]==src and e[:dst]==dst }.first
  end
  
  private

  def resolve_path src, dst, prevs
    prev = prevs[dst]
    return [] if not prev
    path = [ prev, dst ]
    while prev != src
      prev = prevs[prev]
      path.unshift prev
    end
    path
  end

  def spf_result(src, dst, costs, paths)
    nodes   = [dst] if dst
    nodes ||= @nodes
    nodes   = nodes - [src]
    result  = {}
    nodes.each do |node|
      result[node] = {
        'cost' => costs[node],
        'path' => resolve_path(src, node, paths),
      }
    end
    result
  end
end


class LatencyToTopology
  require 'json'
  FILE = 'igp-latency-metric/latency.json'
  attr_reader :nodes, :edges

  def initialize
    data   = JSON.load File.read(FILE)
    @nodes = get_nodes data
    @edges = get_edges data
  end

  def get_nodes data
    data.keys
  end

  def get_edges data
    edges = []
    data.each do |node, peers|
      next unless peers.class == Hash
      peers.each do |peer, peer_data|
        next unless data.class == Hash
        next unless metric = peer_data['latency']
        metric = (metric.first*1000).round
        edges << {
          src: node,
          dst: peer,
          metric: metric,
        }
      end
    end
    edges
  end
end


if $0 == __FILE__
  require 'yaml'
  topo = LatencyToTopology.new
  spf = SPF.new(topo.nodes, topo.edges)
  puts YAML.dump spf.spf ARGV[0], ARGV[1]
  #5.times do
  #  start = Time.now
  #  50.times do
  #    spf.spf ARGV[0], ARGV[1]
  #  end
  #  puts Time.now-start
  #end
  #p spf.cost ARGV[0], ARGV[1]
  #p spf.path ARGV[0], ARGV[1]
end
[ytti@eng0.dllstx09.us.to.gin.ntt.net ~/git/misc]%