Expandable tree with checkboxes

HTML

<script src="https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.17/d3.min.js"></script>

CSS

.node {
  cursor: pointer;
}

.node circle {
  fill: #fff;
  stroke: #C0C0C0;
  stroke-width: 1.5px;
}

.node text {
  font: 10px sans-serif;
}

.link {
  fill: none;
  stroke: #C0C0C0;
  stroke-width: 1.5px;
}

JavaScript

const treeData = {
  "id": 1,
  "name": "Root",
  "checked": false,
  "color": "white",
  "children": [
    {
      "id": 2,
        "name": "Leaf A",
      "checked": false,
      "color": "red",
        "children": [
        {
          "id": 3,
          "name": "A - 1",
                "checked": false,
          "color": "brown",
          }, 
        {
          "id": 4,
          "name": "A - 2",
          "checked": false,
          "color": "orange",
          },
        {
          "id": 5,
          "name": "A - 3",
                "checked": false,
          "color": "yellow",
          },
      ]
    }, 
    {
      "id": 6,
        "name": "Leaf B",
      "checked": false,
      "color": "green",
        "children": [
        {
          "id": 7,
          "name": "B - 1",
          "checked": false,
          "color": "#00ff40",
          }, 
        {
          "id": 8,
          "name": "B - 2",
          "checked": false,
          "color": "#00ff80",
          }
      ]
    }
  ]  
};

const margin = {
    top: 20,
    right: 120,
    bottom: 20,
    left: 120
  };
  
const width = 600 - margin.right - margin.left;
const height = 400 - margin.top - margin.bottom;

var i = 0,duration = 750;
  
const tree = d3.layout.tree()
  .size([height, width]);

const diagonal = d3.svg.diagonal()
  .projection(function(d) {
    return [d.y, d.x];
  });


const svg = d3.select("body").append("svg")
  .attr("width", width + margin.right + margin.left)
  .attr("height", height + margin.top + margin.bottom);
  
const container = svg.append("g")
  .attr("transform", "translate(" + margin.left + "," + margin.top + ")");

const root = treeData;
root.x0 = height / 2;
root.y0 = 0;

// Collapse after the second level
root.children.forEach(collapse);

update(root);

// Collapse the node and all it's children
function collapse(d) {
  if (d.children) {
    d._children = d.children
    d._children.forEach(collapse)
    d.children = null
  }
}

function update(source) {
 ...