
require 'nm'
require 'ksyms'

class Amount
  attr_reader :count, :bytes

  def Amount.Create(x)
    count, bytes = x.unpack("II")
    Amount.new(count, bytes)
  end
  
  def initialize(count=0, bytes=0)
    @count= count
    @bytes = bytes
  end
  
  def + (other)
    Amount.new(@count + other.count, @bytes + other.bytes)
  end
end

class Amount8
  attr_reader :amounts

  def initialize(amounts=nil)
    if (amounts == nil)
      @amounts = [Amount.new]*8
    elsif (amounts.length == 8)
      @amounts = amounts
    else
      abort
    end
  end

  def alloc_amounts
    amounts.slice(0...4)
  end

  def free_amounts
    amounts.slice(4...8)
  end

  def + (other)
    amount8 = Amount8.new
    for i in 0...8
      amount8.amounts[i] = @amounts[i] + other.amounts[i]
    end
    return amount8
  end
end

class Addr_Arc
  attr_reader :amount8, :to, :next, :from

  def Addr_Arc.Create(bin)
    amounts = bin.unpack("IIa8a8a8a8a8a8a8a8")
    to = amounts.shift
    n = amounts.shift
    amount8 = Amount8.new(amounts.collect{|x| Amount.Create(x)})
    Addr_Arc.new(0, to, n, amount8)
  end

  def initialize(from, to, n, amount8)
    @from = from
    @to = to
    @next = n
    @amount8 = amount8
  end

  def fix_from (from)
    @from = from
  end
end

class Sym_Arc
  attr_reader :to_ksym, :from_ksym, :amount8

  def initialize(from, to, amount8)
    @from_ksym = from
    @to_ksym = to
    @amount8 = amount8
  end

  def parallel? (other)
    return (@to_ksym == other.to_ksym && @from_ksym == other.from_ksym)
  end

  def accum(other)
    abort if (!parallel?(other))
    @amount8 += other.amount8
  end
end

class Sym_Node
  attr_accessor :in, :out, :ksym, :in_amount8, :out_amount8

  def initialize (ksym)
    @ksym = ksym
    @in = []
    @out = []
    @in_amount8 = Amount8.new
    @out_amount8 = Amount8.new
  end

  def stick_in (sym_arc)
    @in.push(sym_arc)
    @in_amount8 += sym_arc.amount8
  end

  def stick_out (sym_arc)
    @out.push(sym_arc)
    @out_amount8 += sym_arc.amount8
  end
end

###################
# Display methods #
###################

def display_header (sb, mb, lb)
  printf ("           KMEMPROF ( 0 < s <= %d < m <= %d < l <= %d < x )\n", sb, mb, lb)
  printf ("  Alloc                                    |  Free\n")
  printf ("   bytes   counts     s     m     l     x  |   bytes   counts     s     m     l     x\n")
end

def display_border
  printf ("-------------------------------------------+--------------------------------------------+\n")
end

def display_amounts(amount8)
  alloc = amount8.alloc_amounts
  free = amount8.free_amounts
  [alloc, free].each do |unit|
      bytes = 0
      unit.each { |x| bytes += x.bytes } 
      printf ("%8d ", bytes)
      counts = 0
      unit.each { |x| counts += x.count } 
      printf ("%8d ", counts)
      unit.each { |x| printf ("%5.1f ", counts==0 ? 0 : x.count.to_f*100.0/counts.to_f) }
      printf (" | ")
  end
end

#######################################
# Load kmemprof-cgprof's Profile Data
#######################################

  SIZEOF_INT = 4

  def Load_Bounds (file)
    io = File.open(file)
    header = io.read(7*SIZEOF_INT)
    offset,size,prof_len,prof_shift,small_bound,medium_bound,large_bound=header.unpack("I*")

    return small_bound, medium_bound, large_bound
  end

  def Load_Cgprof_File(ksyms, file)
    io = File.open(file)
    
    # read header
    header = io.read(7*SIZEOF_INT)
    offset,size,prof_len,prof_shift,small_bound,medium_bound,large_bound=header.unpack("I*")
     
    # read table
    table = []
    for i in 0...prof_len
      abort if (ent = io.read(SIZEOF_INT)) == nil
    
      index, = ent.unpack("I")
      if index != 0
        table[index] = (i << prof_shift) + offset
      end
    end
    
    # read buckets
    all_buckets = []
    while  (bucket = io.read(18*SIZEOF_INT))
      all_buckets.push(Addr_Arc.Create(bucket))
    end
    
    # fixate bucket's `from' entry
    buckets = []
    table.each_index do |i|
      addr = table[i]
      next if addr == nil
      while(i != 0)
        bucket = all_buckets[i]
        bucket.fix_from(addr)
        buckets.push(bucket)
        i = bucket.next
      end
    end

    return buckets
  end

#
# bundle Addr_Arcs into Sym_Arc
#

def bundle_addr_arcs (ksyms, addr_arcs)
  tmp_arcs = []
  addr_arcs.each do |addr_arc|
    ksym_from = ksyms.addr2ksym(addr_arc.from)
    ksym_to = ksyms.addr2ksym(addr_arc.to)
    tmp_arcs.push(Sym_Arc.new(ksym_from, ksym_to, addr_arc.amount8))
  end

  bundled_arcs = []
  while tmp_arcs.length > 0
    arc = tmp_arcs.shift
    i = 0
    while i < tmp_arcs.length
      if arc.parallel?(tmp_arcs[i])
        arc.accum(tmp_arcs.delete_at(i))
      else
        i += 1
      end
    end
    bundled_arcs.push(arc)
  end

  return bundled_arcs
end

#
# returns array of Sym_Node
#

def compose_graph(sym_arcs)
  nodes = {}

  sym_arcs.each do |arc|
    to = arc.to_ksym
    from = arc.from_ksym
    next if to == from
    nodes[to] = Sym_Node.new(to) if (nodes[to] == nil)
    nodes[from] = Sym_Node.new(from) if (nodes[from] == nil)
    nodes[to].stick_in(arc)
    nodes[from].stick_out(arc)
  end

  return nodes.values
end

##########
## Main ##
##########

require 'getoptlong'

ksyms_file = '/proc/ksyms'
system_map = '/boot/System.map'
dotty = 0
dotty_block = nil
non_map = 0

opts = GetoptLong.new(
  [ "--ksyms", "-k", GetoptLong::REQUIRED_ARGUMENT],
  [ "--system-map", "-s", GetoptLong::REQUIRED_ARGUMENT],
  [ "--dotty", "-d", GetoptLong::NO_ARGUMENT],
  [ "--dotty-block", "-c", GetoptLong::REQUIRED_ARGUMENT],
  [ "--non-map", "-n", GetoptLong::NO_ARGUMENT]
  )

opts.each do |opt, arg|
  ksyms_file = arg.to_s if opt == "--ksyms"
  system_map = arg.to_s if opt == "--system-map"
  dotty_block = arg.to_s if opt == "--dotty-block"
  dotty = 1 if opt == "--dotty"
  non_map = 1 if opt == "--non-map"
end

ksyms = KSyms.new(ksyms_file, system_map)

addr_arcs = []

small_bound = 0
medium_bound = 0
large_bound = 0

$debug_level = 0

def debug(msg)
  puts msg if $debug_level > 0
end

debug 'loading'

ARGV.each do |arg|
  addr_arcs += Load_Cgprof_File(ksyms, arg.to_s)
  small_bound, medium_bound, large_bound = Load_Bounds(arg.to_s)
end

debug 'loaded'

if (non_map != 0)
  display_header(small_bound, medium_bound, large_bound)
  display_border
  addr_arcs.each do |bucket|
    display_amounts(bucket.amount8)
    printf (" %s -> %s\n", ksyms.addr2symname(bucket.from), 
      ksyms.addr2symname(bucket.to))
    
    display_border
  end
  exit
end

debug 'bandle..'
bundled_arcs = bundle_addr_arcs(ksyms, addr_arcs)
debug 'bundled'

def dotty_output (sym_arcs)
  puts "digraph G {"

  sym_arcs.each do |arc|
    label = nil
    if block_given?
      label = yield arc.amount8
    else
      total=0
      arc.amount8.alloc_amounts.each {|x| total+=x.count}
      label = total.to_s if total > 0
    end
    if label != nil
      printf ("\"%s\" -> \"%s\" [label=\"%s\"]\n",
        arc.from_ksym.nm_sym.name, arc.to_ksym.nm_sym.name, label)
    end
  end

  puts "}"
end

if (dotty != 0)
  dotty_output(bundled_arcs)
end

debug 'compose..'
graph_nodes = compose_graph(bundled_arcs)
debug 'composed'

### display

display_header(small_bound, medium_bound, large_bound)
display_border

# kick out 
def sym_exp(ksym)
  sym = ksym.nm_sym.name
  sym += '[' + ksym.mod_name + ']' if (ksym.mod_name != '')
  return sym
end

graph_nodes.each do |node|
  node.in.each do |arc|
    display_amounts( arc.amount8 )
    puts "    " + sym_exp(arc.from_ksym)
  end

  total_amounts = (node.in.length == 0) ? node.out_amount8 : node.in_amount8

  display_amounts(total_amounts)
  puts sym_exp(node.ksym)

  node.out.each do |arc|
    display_amounts( arc.amount8 )
    puts "    " + sym_exp(arc.to_ksym)
  end

  display_border
end

