quadtree lobster

quadtree_visualizer.lobster

import vec
import color
import gl
import std

let WIDTH = 800.0
let HEIGHT = 600.0

class Circle:
    position: float2
    radius: float
    color: color
    original_color: color

struct Rect:
    x: float
    y: float
    w: float
    h: float

class Quadtree:
    boundary: Rect
    capacity: int
    circles: [Circle] = []
    divided = false
    northeast: Quadtree?
    northwest: Quadtree?
    southeast: Quadtree?
    southwest: Quadtree?

def rect_contains_circle(r::Rect, c::Circle) -> bool:
    return c.position.x >= r.x and c.position.x < r.x + r.w and
           c.position.y >= r.y and c.position.y < r.y + r.h

def rect_intersects_rect(r1:Rect, r2:Rect) -> bool:
    return not (r2.x >= r1.x + r1.w or
                r2.x + r2.w <= r1.x or
                r2.y >= r1.y + r1.h or
                r2.y + r2.h <= r1.y)

def subdivide(qt::Quadtree):
    let x = qt.boundary.x
    let y = qt.boundary.y
    let hw = qt.boundary.w / 2.0
    let hh = qt.boundary.h / 2.0

    qt.northeast = Quadtree { Rect { x + hw, y,      hw, hh }, qt.capacity, northeast: nil, northwest: nil, southeast: nil, southwest: nil}
    qt.northwest = Quadtree { Rect { x,      y,      hw, hh }, qt.capacity, northeast: nil, northwest: nil, southeast: nil, southwest: nil}
    qt.southeast = Quadtree { Rect { x + hw, y + hh, hw, hh }, qt.capacity, northeast: nil, northwest: nil, southeast: nil, southwest: nil}
    qt.southwest = Quadtree { Rect { x,      y + hh, hw, hh }, qt.capacity, northeast: nil, northwest: nil, southeast: nil, southwest: nil}
    qt.divided = true

def insertik(qt::Quadtree, c::Circle) -> bool:
    if not rect_contains_circle(qt.boundary, c):
        return false

    if qt.circles.length < qt.capacity:
        qt.circles.push(c)
        return true
    else:
        if not qt.divided:
            qt.subdivide()

        if qt.northeast and qt.northeast.insertik(c): return true
        if qt.northwest and qt.northwest.insertik(c): return true
        if qt.southeast and qt.southeast.insertik(c): return true
        if qt.southwest and qt.southwest.insertik(c): return true
        
        qt.circles.push(c)
        return true

def query(qt::Quadtree, range_rect::Rect, found_circles:[Circle]) -> void:
    if not rect_intersects_rect(qt.boundary, range_rect):
        return

    for (qt.circles) c_item:
        let circle_rect = Rect { c_item.position.x - c_item.radius, 
                                 c_item.position.y - c_item.radius, 
                                 c_item.radius * 2.0, 
                                 c_item.radius * 2.0 }
        if rect_intersects_rect(range_rect, circle_rect):
            found_circles.push(c_item)

    if qt.divided:
        if qt.northeast: qt.northeast.query(range_rect, found_circles)
        if qt.northwest: qt.northwest.query(range_rect, found_circles)
        if qt.southeast: qt.southeast.query(range_rect, found_circles)
        if qt.southwest: qt.southwest.query(range_rect, found_circles)

def draw_quadtree_boundaries(qt::Quadtree) -> void:
    gl.line_mode(1)
    gl.color(color_blue)
    gl.translate(float2 { qt.boundary.x, qt.boundary.y }):
        gl.rect(float2 { qt.boundary.w, qt.boundary.h })
    gl.line_mode(0)

    if qt.divided:
        if qt.northeast: draw_quadtree_boundaries(qt.northeast)
        if qt.northwest: draw_quadtree_boundaries(qt.northwest)
        if qt.southeast: draw_quadtree_boundaries(qt.southeast)
        if qt.southwest: draw_quadtree_boundaries(qt.southwest)

fatal(gl.window("Quadtree Example with Circles", int(WIDTH), int(HEIGHT)))

let num_circles = 1000
let circles = map(num_circles):
    let x = rnd_float() * WIDTH
    let y = rnd_float() * HEIGHT
    let radius = 5.0
    Circle { float2 { x, y }, radius, color_red, color_red }

let boundary = Rect { 0.0, 0.0, WIDTH, HEIGHT }
let qtree = Quadtree { boundary, 4, northeast: nil, northwest: nil, southeast: nil, southwest: nil} 

for (circles) c:
    qtree.insertik(c)

var show_quadtree_lines = false

while gl.frame() and gl.button("escape") != 1:
    gl.clear(color_white)

    let mouse_pixel_pos = gl.mouse_pos(0)
    let mouse_pos = float(mouse_pixel_pos)
    let query_range_size = 100.0
    let mouse_range = Rect {
        mouse_pos.x - query_range_size / 2.0,
        mouse_pos.y - query_range_size / 2.0,
        query_range_size,
        query_range_size
    }

    let found_circles_in_range = []
    qtree.query(mouse_range, found_circles_in_range)

    for (circles) c:
        c.color = c.original_color
    
    for (found_circles_in_range) c:
        c.color = color_green

    for (circles) c:
        gl.color(c.color)
        gl.translate(c.position):
            gl.circle(c.radius, 20)

    gl.color(color { 0.0, 0.0, 0.8, 0.3 })
    gl.translate(float2 { mouse_range.x, mouse_range.y }):
        gl.rect(float2 { mouse_range.w, mouse_range.h })

    if gl.button("q") == 1: show_quadtree_lines = not show_quadtree_lines
    if show_quadtree_lines:
        draw_quadtree_boundaries(qtree)