desk2/client/bvh.js

function bvh_make_leaf(bvh, index, stroke) {
    const leaf = {
        'stroke_index': index,
        'bbox': stroke.bbox,
        'area': stroke.area,
        'parent_index': null,
        'is_leaf': true,
    };

    bvh.nodes.push(leaf);

    return bvh.nodes.length - 1;
}

function bvh_make_internal(bvh) {
    const node = {
        'child1': null,
        'child2': null,
        'parent_index': null,
        'is_leaf': false,
    };

    bvh.nodes.push(node);

    return bvh.nodes.length - 1;
}

function bvh_compute_sah(bvh, new_leaf, potential_sibling, only_parent = false) {
    let cost = 0;
    let union_box;

    if (!only_parent) {
        union_box = quad_union(new_leaf.bbox, potential_sibling.bbox);

        const internal_node_would_be = { 'bbox': union_box };
        const new_internal_node_cost = (union_box.x2 - union_box.x1) * (union_box.y2 - union_box.y1);

        cost += new_internal_node_cost;
    } else {
        union_box = new_leaf.bbox;
    }

    let parent_index = potential_sibling.parent_index;

    while (parent_index !== null) {
        const current_node = bvh.nodes[parent_index];
        const old_cost = current_node.area; 
        union_box = quad_union(current_node.bbox, union_box); 
        const new_cost = (union_box.x2 - union_box.x1) * (union_box.y2 - union_box.y1);
        cost += new_cost - old_cost;
        parent_index = current_node.parent_index;
    }

    return cost;
}

function bvh_find_best_sibling(bvh, leaf_index) {
    // branch and bound

    const leaf = bvh.nodes[leaf_index];
    const leaf_cost = (leaf.bbox.x2 - leaf.bbox.x1) * (leaf.bbox.y2 - leaf.bbox.y1);

    let best_cost = bvh_compute_sah(bvh, leaf, bvh.nodes[bvh.root]);
    let best_index = bvh.root;

    bvh.pqueue.clear();
    bvh.pqueue.push(best_index, best_cost);

    while (bvh.pqueue.size > 0) {
        const current_index = bvh.pqueue.pop();
        const current_node = bvh.nodes[current_index];
        const cost = bvh_compute_sah(bvh, current_node, leaf);

        if (cost < best_cost) {
            best_cost = cost;
            best_index = current_index;
        }

        if (!current_node.is_leaf) {
            const child1 = bvh.nodes[current_node.child1];
            const lower_bound_for_children = bvh_compute_sah(bvh, child1, leaf, true) + leaf_cost;
            if (lower_bound_for_children < best_cost) {
                bvh.pqueue.push(current_node.child1, lower_bound_for_children);
                bvh.pqueue.push(current_node.child2, lower_bound_for_children);
            }
        }
    }

    return best_index;
}

function bvh_add_stroke(bvh, index, stroke) {
    const leaf_index = bvh_make_leaf(bvh, index, stroke);

    if (bvh.nodes.length === 1) {
        bvh.root = leaf_index;
        return;
    }

    if (bvh.pqueue.capacity < Math.ceil(bvh.nodes.length * 1.2)) {
        bvh.pqueue = new MinQueue(bvh.nodes.length * 2);        
    }

    // It's as easy as 1-2-3

    // 1. Find best sibling for leaf
    const sibling = bvh_find_best_sibling(bvh, leaf_index);

    // 2. Create new parent
    const old_parent = bvh.nodes[sibling].parent_index;
    const new_parent = bvh_make_internal(bvh);

    bvh.nodes[new_parent].parent_index = old_parent;
    bvh.nodes[new_parent].bbox = quad_union(stroke.bbox, bvh.nodes[sibling].bbox);

    if (old_parent !== null) {
        // The sibling was not the root
        if (bvh.nodes[old_parent].child1 === sibling) {
            bvh.nodes[old_parent].child1 = new_parent;
        } else {
            bvh.nodes[old_parent].child2 = new_parent;
        }
       
        bvh.nodes[new_parent].child1 = sibling;
        bvh.nodes[new_parent].child2 = leaf_index;
        
        bvh.nodes[sibling].parent_index = new_parent;
        bvh.nodes[leaf_index].parent_index = new_parent;
    } else {
        // The sibling was the root
        bvh.nodes[new_parent].child1 = sibling;
        bvh.nodes[new_parent].child2 = leaf_index;

        bvh.nodes[sibling].parent_index = new_parent;
        bvh.nodes[leaf_index].parent_index = new_parent;

        bvh.root = new_parent;
    }

    const new_bbox = quad_union(bvh.nodes[bvh.nodes[new_parent].child1].bbox, bvh.nodes[bvh.nodes[new_parent].child2].bbox);
    bvh.nodes[new_parent].bbox = new_bbox;
    bvh.nodes[new_parent].area = (new_bbox.x2 - new_bbox.x1) * (new_bbox.y2 - new_bbox.y1);   

    // 3. Refit and insert in fullnode 
    let refit_index = bvh.nodes[leaf_index].parent_index;
    while (refit_index !== null) {
        const child1 = bvh.nodes[refit_index].child1;
        const child2 = bvh.nodes[refit_index].child2;

        bvh.nodes[refit_index].bbox = quad_union(bvh.nodes[child1].bbox, bvh.nodes[child2].bbox);

        if (bvh.nodes[refit_index].is_fullnode) {
            tv_add2(bvh.nodes[refit_index].stroke_indices, index); 
        }

        refit_index = bvh.nodes[refit_index].parent_index;
    }
}

function bvh_intersect_quad(bvh, quad, result_buffer) {
    if (bvh.root === null) {
        return;
    }
    
    tv_clear(bvh.traverse_stack);
    tv_add(bvh.traverse_stack, bvh.root);

    while (bvh.traverse_stack.size > 0) {
        const node_index = tv_pop(bvh.traverse_stack);
        const node = bvh.nodes[node_index];

        if (!quads_intersect(node.bbox, quad)) {
            continue;
        }

        if (node.is_fullnode) {
            if (quad_fully_inside(quad, node.bbox)) {
                tv_append(result_buffer, tv_data(node.stroke_indices)); 
                continue;
            }
        }

        if (node.is_leaf) {
            tv_add(result_buffer, node.stroke_index);
        } else {
            tv_add(bvh.traverse_stack, node.child1);
            tv_add(bvh.traverse_stack, node.child2);
        }
    }
}

function bvh_clip(state, context) {
    if (state.stroke_count === 0) {
        return;
    }

    tv_ensure(context.clipped_indices, round_to_pow2(state.stroke_count, 4096))
    tv_ensure(state.bvh.traverse_stack, round_to_pow2(state.stroke_count, 4096));

    tv_clear(context.clipped_indices);

    const screen_topleft     = screen_to_canvas(state, {'x': 0, 'y': 0});
    const screen_bottomright = screen_to_canvas(state, {'x': context.canvas.width, 'y': context.canvas.height});
    const screen_topright    = { 'x': screen_bottomright.x, 'y': screen_topleft.y };
    const screen_bottomleft  = { 'x': screen_topleft.x,     'y': screen_bottomright.y };

    const screen = {
        'x1': screen_topleft.x,
        'y1': screen_topleft.y,
        'x2': screen_bottomright.x,
        'y2': screen_bottomright.y
    };

    bvh_intersect_quad(state.bvh, screen, context.clipped_indices);

    tv_data(context.clipped_indices).sort(); // we need to draw back to front still!
}

function bvh_point(state, p) {
    const bvh = state.bvh;
    const stack = [];
    const indices = [];

    if (bvh.root === null) {
        return null;
    }
    
    stack.push(bvh.root);

    while (stack.length > 0) {
        const node_index = stack.pop();
        const node = bvh.nodes[node_index];

        if (!point_in_bbox(p, node.bbox)) {
            continue;
        }

        if (node.is_leaf) {
            const stroke = state.events[node.stroke_index];
            const xs = state.wasm.buffers['xs'].tv.data.subarray(stroke.coords_from, stroke.coords_to);
            const ys = state.wasm.buffers['ys'].tv.data.subarray(stroke.coords_from, stroke.coords_to);
            const pressures = state.wasm.buffers['pressures'].tv.data.subarray(stroke.coords_from, stroke.coords_to);

            if (point_in_stroke(p, xs, ys, pressures, stroke.width)) {
                indices.push(node.stroke_index);
            }
        } else {
            stack.push(node.child1);
            stack.push(node.child2);
        }
    }

    if (indices.length > 0) {
        indices.sort();
        return indices[indices.length - 1];
    }

    return null;
}

function bvh_construct_rec(bvh, vertical, strokes, depth) {
    if (strokes.length > 1) {
        // internal
        let sorted_strokes;

        if (vertical) {
            sorted_strokes = strokes.toSorted((a, b) => a.bbox.cy - b.bbox.cy);
        } else {
            sorted_strokes = strokes.toSorted((a, b) => a.bbox.cx - b.bbox.cx);
        }

        const node_index = bvh_make_internal(bvh);
        const left_of_split_count = Math.floor(strokes.length / 2);

        const child1 = bvh_construct_rec(bvh, !vertical, sorted_strokes.slice(0, left_of_split_count), depth + 1);
        const child2 = bvh_construct_rec(bvh, !vertical, sorted_strokes.slice(left_of_split_count, sorted_strokes.length), depth + 1);

        bvh.nodes[child1].parent_index = node_index; 
        bvh.nodes[child2].parent_index = node_index;

        bvh.nodes[node_index].child1 = child1;
        bvh.nodes[node_index].child2 = child2;
        bvh.nodes[node_index].bbox = quad_union(bvh.nodes[child1].bbox, bvh.nodes[child2].bbox);

        if (depth === config.bvh_fullnode_depth) {
            const indices = tv_create(Int32Array, round_to_pow2(strokes.length, 32));
            
            for (let i = 0; i < strokes.length; ++i) {
                tv_add(indices, strokes[i].index);
            }

            bvh.nodes[node_index].stroke_indices = indices;
            bvh.nodes[node_index].is_fullnode = true;
        }

        return node_index;
    } else {
        // leaf
        return bvh_make_leaf(bvh, strokes[0].index, strokes[0]);
    }
}

function bvh_construct(state) {
    if (state.events.length > 0) {
        state.bvh.root = bvh_construct_rec(state.bvh, true, state.events, 0); 
    }
}