#!/usr/bin/env ruby
#
# Copyright © 2015 Siarhei Siamashka <siarhei.siamashka@gmail.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

###############################################################################
# This is a u-boot SPL stack usage analysis script. Needs to be run           #
# from the root of the u-boot source tree after a successful SPL build.       #
#                                                                             #
# Produces 'spl-stackgraph.png' file as the output.                           #
###############################################################################

require 'find'

###############################################################################

if !File.directory?("spl")
  abort("Error: can't find the 'spl' directory.\n" +
        "Please ensure that you are in the root of the u-boot source tree.\n")
end

def tool_exists(tool_name)
    `which #{tool_name} > /dev/null 2>&1`
    return $?.to_i == 0
end

if !tool_exists("dot")
  abort("Please install the graphviz dot tool.\n")
end

###############################################################################
# Get objdump log                                                             #
###############################################################################

# Get the information from the environment
$toolchain = ENV["CROSS_COMPILE"]

# Try to guess the ARM toolchain automatically
if !$toolchain && `file spl/u-boot-spl` =~ / ARM, /
  toolchain_list = [
    "arm-none-linux-gnueabi-",
    "armv7a-hardfloat-linux-gnueabi-",
    "arm-linux-gnueabihf-"
  ]
  toolchain_list.each {|toolchain|
    if tool_exists("#{toolchain}objdump")
      $toolchain = toolchain
      break
    end
  }
end

if !$toolchain
  abort("Error: can't find a usable toolchain.\n" +
        "Try to set the CROSS_COMPILE environment variable.\n")
end

$objdump_dis_log = `#{$toolchain}objdump -d spl/u-boot-spl`
$objdump_sym_log = `#{$toolchain}objdump -t spl/u-boot-spl`
$size_log = `#{$toolchain}size spl/u-boot-spl`

###############################################################################
# Read the *.su files (generated by GCC with '-fstack-usage' option)          #
###############################################################################

$self_stack_size = {}
$total_stack_size  = {}

Find.find("spl") { |filename|
  next unless filename =~ /.*\.su$/
  File.open(filename).each_line {|l|
    if l =~ /\:([^\:\s]+)\s+(\d+)\s+(\w+)$/
      $self_stack_size[$1] = [$2.to_i, ($self_stack_size[$1] || 0)].max
    end
  }
}

if $self_stack_size.size == 0
  abort("Error: could not load anything from *.su files in 'spl' directory.\n" +
        "Please ensure that you are using '-fstack-usage' GCC option.\n")
end

###############################################################################
# Parse the objdump log                                                       #
###############################################################################

$detected_cycles = {}
$labels = {}
$edges = {}
$functions = {}
$refs = {}
$xrefs = {}
$max_stack_size = 0
$critical_path = {}

$objdump_sym_log.each_line {|l|
  $functions[$1] = 1 if l =~ /F\s+\.text.*?(\S+)$/
}

current_func = "entry"
$objdump_dis_log.each_line {|l|
  current_func = $1 if l =~ /^\h+\s+\<(.*?)>\:$/ && $functions.has_key?($1)
  if l =~ /(.*?)<(.*?)(\+0x\h+)?>/ && $2 != current_func
    cmd = $1
    func_name = $2
    $edges[[current_func, func_name]] = 1
    $refs[current_func] = {} if !$refs.has_key?(current_func)
    $refs[func_name] = {} if !$refs.has_key?(func_name)
    $refs[current_func][func_name] = ($self_stack_size[func_name] || 0)

    $xrefs[current_func] = {} if !$xrefs.has_key?(current_func)
    $xrefs[func_name] = {} if !$xrefs.has_key?(func_name)
    $xrefs[func_name][current_func] = ($self_stack_size[current_func] || 0)
  end
}

def recursive_walk(node, visited_nodes, total_stack)
  $total_stack_size[node] = [total_stack, ($total_stack_size[node] || 0)].max
  if $total_stack_size[node] > $max_stack_size
    $critical_path = visited_nodes.clone
    $max_stack_size = $total_stack_size[node]
  end
  $refs[node].each {|next_node, func_stack|
    if visited_nodes.has_key?(next_node) then
      $detected_cycles[[node, next_node]] = 1
      next
    end
    visited_nodes[next_node] = 1
    recursive_walk(next_node, visited_nodes, total_stack + func_stack + 8)
    visited_nodes.delete(next_node)
  }
end

# Find the 'root' nodes and recursively traverse the graph starting from them
$xrefs.each {|func_name, nodes|
  next if nodes.size != 0
  visited_nodes = {}
  visited_nodes[func_name] = 1
  recursive_walk(func_name, visited_nodes, ($self_stack_size[func_name] || 0))
}

###############################################################################
# Generate .dot file for graphviz                                             #
###############################################################################

$legend = ""
$max_stack_size = $total_stack_size.map { |k, v| v }.max

if $size_log =~ /(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\h+)\s+spl\/u\-boot\-spl/
    text_plus_data_size = $1.to_i + $2.to_i
    bss_size = $3.to_i

    $legend << sprintf("------------ u-boot SPL ------------\\l")
    $legend << sprintf("text+data   = %5d bytes\\l", text_plus_data_size)
    $legend << sprintf("bss         = %5d bytes\\l", bss_size)
    $legend << sprintf("stack usage = %5d bytes\\l", $max_stack_size)
end

if $detected_cycles.size != 0
  $legend << sprintf("--------------------------------------\\l")
  $legend << sprintf("Warning: cycles detected in functions:\\l")
  tmp = {}
  $detected_cycles.each {|k, v|
    tmp[k[0]] = 1
    tmp[k[1]] = 1
  }
  tmp.each {|k, v| $legend << sprintf("    %s\\l", k) }
  $legend << sprintf("Watch for bold red arrows on\\l")
  $legend << sprintf("the call graph picture.\\l")
end

fh = File.open("spl-stackgraph.dot", "w")

$refs.each {|func_name, nodes|
  $labels[func_name] = sprintf("%s\\n[total=%d, self=%d]",
                               func_name,
                               ($total_stack_size[func_name] || 0),
                               ($self_stack_size[func_name] || 0))
}

fh.printf("digraph graphname {\n")
fh.printf("rankdir=LR;\n")
fh.printf("node [shape=box, fontname=arial];\n")

$total_stack_size.each { |k, v|
  next if !$critical_path.has_key?(k)
  fh.printf("\"%s\" [style=filled,color=orange];\n", $labels[k])
}

$edges.each {|k, v|
  if $detected_cycles.has_key?(k) then
    fh.printf("\"%s\" -> \"%s\" [color=\"red\",penwidth=5];\n",
              $labels[k[0]], $labels[k[1]])
  else
    fh.printf("\"%s\" -> \"%s\";\n", $labels[k[0]], $labels[k[1]])
  end
}
fh.printf("legend [shape=box,style=filled,fontname=monospace,fontsize=24," +
          "color=\"lightgray\",label=\"#{$legend}\"];\n")
fh.printf("}\n")
fh.close

`dot -Tpng -o spl-stackgraph.png spl-stackgraph.dot`