#!/usr/bin/env dart
// Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

import "dart:io";

class Symbol {
  String? name;
  String? type;
  int? shallowSize;
  int? retainedSize;
  List<Symbol> children = <Symbol>[];

  Symbol compressTrivialPaths() {
    for (var i = 0; i < children.length; i++) {
      children[i] = children[i].compressTrivialPaths();
    }
    if ((type == "path") && (children.length == 1)) {
      return children[0];
    }
    return this;
  }

  int computeRetainedSize() {
    var s = shallowSize!;
    for (var child in children) {
      s += child.computeRetainedSize();
    }
    retainedSize = s;
    return s;
  }

  writeJson(StringBuffer out) {
    out.write("{");
    out.write('"name": "$name",');
    out.write('"shallowSize": $shallowSize,');
    out.write('"retainedSize": $retainedSize,');
    out.write('"type": "$type",');
    out.write('"children": [');
    bool first = true;
    for (var child in children) {
      if (first) {
        first = false;
      } else {
        out.write(",");
      }
      child.writeJson(out);
    }
    out.write("]}");
  }
}

const filteredPathComponents = <String>[
  "",
  ".",
  "..",
  "out",
  "xcodebuild",
  "src",
  "source",
  "third_party",
  "DebugX64",
  "ReleaseX64",
  "ProductX64",
  "DebugARM64",
  "ReleaseARM64",
  "ProductARM64",
  "DebugXARM64",
  "ReleaseXARM64",
  "ProductXARM64",
];

var cwd = Directory.current.path;
String prettyPath(String path) {
  if (path.startsWith(cwd)) {
    path = path.substring(cwd.length);
  }
  return path
      .split("/")
      .where((component) => !filteredPathComponents.contains(component))
      .join("/");
}

main(List<String> args) {
  if (args.isEmpty) {
    print("Usage: binary_size <binaries>");
    exit(1);
  }

  for (var arg in args) {
    analyze(arg);
  }
}

analyze(String binaryPath) {
  var nmExec = "nm";
  var nmArgs = ["-a", "--demangle", "--line-numbers",
                "--print-size", binaryPath];
  var nmResult = Process.runSync(nmExec, nmArgs);
  if (nmResult.exitCode != 0) {
    print("+ ${nmExec} ${nmArgs.join(' ')}");
    print(nmResult.exitCode);
    print(nmResult.stdout);
    print(nmResult.stderr);
    exit(1);
  }

  var root = new Symbol();
  root.name = binaryPath;
  root.type = "path";
  root.shallowSize = 0;
  var paths = new Map<String, Symbol>();
  paths[""] = root;
  addToPath(Symbol s, String path) {
    Symbol? p = paths[path];
    if (p == null) {
      p = new Symbol();
      p.name = path;
      p.type = "path";
      p.shallowSize = 0;
      paths[path] = p;

      var i = path.lastIndexOf("/");
      if (i != -1) {
        p.name = path.substring(i + 1);
        addToPath(p, path.substring(0, i));
      } else {
        root.children.add(p);
      }
    }
    p.children.add(s);
  }

  var lines = nmResult.stdout.split("\n");
  var regex = new RegExp("([0-9a-f]+) ([0-9a-f]+) ([a-zA-z]) (.*)");
  for (var line in lines) {
    print(line);
    var match = regex.firstMatch(line);
    if (match == null) {
      continue;
    }

    var address = int.parse(match[1]!, radix: 16);
    var size = int.parse(match[2]!, radix: 16);
    var type = match[3];
    if (type == "b" || type == "B") {
      // Uninitialized data does not contribute to file size.
      continue;
    }

    var nameAndPath = match[4]!.split("\t");
    var name = nameAndPath[0].trim();
    var path = nameAndPath.length == 1 ? "" : nameAndPath[1].trim();
    var colon = path.lastIndexOf(":");
    if (colon > 0) {
      path = path.substring(0, colon);
    }
    path = prettyPath(path);

    for (var prefix in const ["vtable for ",
                              "typeinfo for ",
                              "typeinfo name for "]) {
      if (name.startsWith(prefix)) {
        name = name.substring(prefix.length);
        path = prefix;
      }
    }

    var s = new Symbol();
    s.name = name;
    s.type = type;
    s.shallowSize = size;
    addToPath(s, path);
  }

  root.compressTrivialPaths();
  root.computeRetainedSize();

  var json = new StringBuffer();
  root.writeJson(json);

  var html = viewer.replaceAll("__DATA__", json.toString());
  new File("${binaryPath}.html").writeAsStringSync(html);

  // This written as a URL instead of path because some terminals will
  // automatically recognize it and make it a link.
  var url = Directory.current.uri.resolve("${binaryPath}.html");
  print("Wrote $url");
}

var viewer = """
<html>
<head>
<style>
.treemapTile {
    position: absolute;
    box-sizing: border-box;
    border: solid 1px;
    font-size: 13;
    text-align: center;
    overflow: hidden;
    white-space: nowrap;
    cursor: default;
}
</style>
</head>
<body>
<script>
var root = __DATA__;

function hash(string) {
  // Jenkin's one_at_a_time.
  let h = string.length;
  for (let i = 0; i < string.length; i++) {
    h += string.charCodeAt(i);
    h += h << 10;
    h ^= h >> 6;
  }
  h += h << 3;
  h ^= h >> 11;
  h += h << 15;
  return h;
}

function color(string) {
  let hue = hash(string) % 360;
  return "hsl(" + hue + ",90%,80%)";
}

function prettySize(size) {
  if (size < 1024) return size + "B";
  size /= 1024;
  if (size < 1024) return size.toFixed(1) + "KiB";
  size /= 1024;
  if (size < 1024) return size.toFixed(1) + "MiB";
  size /= 1024;
  return size.toFixed(1) + "GiB";
}

function createTreemapTile(v, width, height, depth) {
  let div = document.createElement("div");
  div.className = "treemapTile";
  div.style["background-color"] = color(v.type);
  div.ondblclick = function(event) {
    event.stopPropagation();
    if (depth == 0) {
      let parent = v.parent;
      if (parent === undefined) {
        // Already at root.
      } else {
        showDominatorTree(parent);  // Zoom out.
      }
    } else {
      showDominatorTree(v);  // Zoom in.
    }
  };

  let left = 0;
  let top = 0;

  const kPadding = 5;
  const kBorder = 1;
  left += kPadding - kBorder;
  top += kPadding - kBorder;
  width -= 2 * kPadding;
  height -= 2 * kPadding;

  div.title =
    v.name +
    " \\ntype: " + v.type +
    " \\nretained: " + v.retainedSize +
    " \\nshallow: " + v.shallowSize;

  if (width < 10 || height < 10) {
    // Too small: don't render label or children.
    return div;
  }

  let label = v.name + " [" + prettySize(v.retainedSize) + "]";
  div.appendChild(document.createTextNode(label));
  const kLabelHeight = 13;
  top += kLabelHeight;
  height -= kLabelHeight;

  if (depth > 2) {
    // Too deep: don't render children.
    return div;
  }
  if (width < 4 || height < 4) {
    // Too small: don't render children.
    return div;
  }

  let children = new Array();
  v.children.forEach(function(c) {
    // Size 0 children seem to confuse the layout algorithm (accumulating
    // rounding errors?).
    if (c.retainedSize > 0) {
      children.push(c);
    }
  });
  children.sort(function (a, b) {
    return b.retainedSize - a.retainedSize;
  });

  const scale = width * height / v.retainedSize;

  // Bruls M., Huizing K., van Wijk J.J. (2000) Squarified Treemaps. In: de
  // Leeuw W.C., van Liere R. (eds) Data Visualization 2000. Eurographics.
  // Springer, Vienna.
  for (let rowStart = 0;  // Index of first child in the next row.
       rowStart < children.length;) {
    // Prefer wider rectangles, the better to fit text labels.
    const GOLDEN_RATIO = 1.61803398875;
    let verticalSplit = (width / height) > GOLDEN_RATIO;

    let space;
    if (verticalSplit) {
      space = height;
    } else {
      space = width;
    }

    let rowMin = children[rowStart].retainedSize * scale;
    let rowMax = rowMin;
    let rowSum = 0;
    let lastRatio = 0;

    let rowEnd;  // One after index of last child in the next row.
    for (rowEnd = rowStart; rowEnd < children.length; rowEnd++) {
      let size = children[rowEnd].retainedSize * scale;
      if (size < rowMin) rowMin = size;
      if (size > rowMax) rowMax = size;
      rowSum += size;

      let ratio = Math.max((space * space * rowMax) / (rowSum * rowSum),
                           (rowSum * rowSum) / (space * space * rowMin));
      if ((lastRatio != 0) && (ratio > lastRatio)) {
        // Adding the next child makes the aspect ratios worse: remove it and
        // add the row.
        rowSum -= size;
        break;
      }
      lastRatio = ratio;
    }

    let rowLeft = left;
    let rowTop = top;
    let rowSpace = rowSum / space;

    for (let i = rowStart; i < rowEnd; i++) {
      let child = children[i];
      let size = child.retainedSize * scale;

      let childWidth;
      let childHeight;
      if (verticalSplit) {
        childWidth = rowSpace;
        childHeight = size / childWidth;
      } else {
        childHeight = rowSpace;
        childWidth = size / childHeight;
      }

      let childDiv = createTreemapTile(child, childWidth, childHeight, depth + 1);
      childDiv.style.left = rowLeft + "px";
      childDiv.style.top = rowTop + "px";
      // Oversize the final div by kBorder to make the borders overlap.
      childDiv.style.width = (childWidth + kBorder) + "px";
      childDiv.style.height = (childHeight + kBorder) + "px";
      div.appendChild(childDiv);

      if (verticalSplit)
        rowTop += childHeight;
      else
        rowLeft += childWidth;
    }

    if (verticalSplit) {
      left += rowSpace;
      width -= rowSpace;
    } else {
      top += rowSpace;
      height -= rowSpace;
    }

    rowStart = rowEnd;
  }

  return div;
}

function showDominatorTree(v) {
  // Add the content div to the document first so the browser will calculate
  // the available width and height.
  let w = document.body.offsetWidth;
  let h = document.body.offsetHeight;

  let topTile = createTreemapTile(v, w, h, 0);
  topTile.style.width = w;
  topTile.style.height = h;
  setBody(topTile);
}

function setBody(div) {
  let body = document.body;
  while (body.firstChild) {
    body.removeChild(body.firstChild);
  }
  body.appendChild(div);
}

function setParents(v) {
  v.children.forEach(function (child) {
    child.parent = v;
    setParents(child);
  });
}

setParents(root);
showDominatorTree(root);

</script>
</body>
</html>
""";