Feedback on Spatial Hash Grid

Help and Feedback Code Review
#1

Hey everyone!

I’ve implemented a spatial hash grid for efficiently searching for nearby objects.
I’d greatly appreciate any advice on changes on how to improve readabillity or performance.

Heres a basic explanation of how it works:

  • The space is divided into a grid of fixed-size cells.
  • Objects are hashed into cells based on their position.
  • For any radius search, you only need to search the relevant cells and their neighbors rather than the entire grid.

Any feedback would be appreciated, Thanks!

--!native
--!strict

export type Grid = typeof(setmetatable({} :: {
    _cellSize: number,
    _objectCellMap: { [any]: CellId },
    _cells: { [CellId]: Cell}
}, {} :: GridImpl))

type GridImpl = {
    __index: GridImpl,

    new: (cellSize: number) -> Grid,
    Insert: (self: Grid, object: any, position: Vector3) -> (),
    Update: (self: Grid, object: any, position: Vector3) -> (),
    Remove: (self: Grid, object: any) -> (),
    RadiusSearch: (self: Grid, from: Vector3, radius: number) -> ({ any }, { number }),
    RadiusSearchAmount: (self: Grid, from: Vector3, radius: number, amount: number) -> { any }
}

type CellId = string

type Cell = { [any]: Vector3 }

local function cellPositionToCellId(positionX: number, positionZ: number)
    return positionX .. positionZ
end

local function positionToCellPosition(position: Vector3, cellSize: number): (number, number)
    return math.floor(position.X / cellSize), math.floor(position.Z / cellSize)
end

local Grid = {} :: GridImpl
Grid.__index = Grid

function Grid.new(cellSize: number)
    return setmetatable({
        _cellSize = cellSize,
        _objectCellMap = {},
        _cells = {}
    }, Grid)
end

function Grid:Insert(object: any, position: Vector3)
    assert(object, "Object cannot be of type nil")
    assert(typeof(position) == "Vector3", "Position is not of type Vector3")
    assert(not self._objectCellMap[object], "Object already exists in the grid, did you mean to call Grid:Update() ?")

    local cellPositionX, cellPositionZ = positionToCellPosition(position, self._cellSize)
    local cellId = cellPositionToCellId(cellPositionX, cellPositionZ)

    self._objectCellMap[object] = cellId
    
    if self._cells[cellId] then
        self._cells[cellId][position] = object
        return
    end

    local cell = {
        [object] = position
    }

    self._cells[cellId] = cell
end

function Grid:Remove(object: any)
    assert(object, "Object cannot be of type nil")
    assert(self._objectCellMap[object], "Object does not exist in the grid, did you mean to call Grid:Insert() ?")

    local cellId = self._objectCellMap[object]
    self._cells[cellId][object] = nil
    self._objectCellMap[object] = nil
end

function Grid:Update(object: any, position: Vector3)
    assert(object, "Object cannot be of type nil")
    assert(typeof(position) == "Vector3", "Position is not of type Vector3")

    -- if the object hasnt moved cells, we should do it here
    local cellPositionX, cellPositionZ = positionToCellPosition(position, self._cellSize)
    local newCellId = cellPositionToCellId(cellPositionX, cellPositionZ)

    if self._objectCellMap[object] == newCellId then
        self._cells[newCellId][object] = position
        return
    end

    self:Remove(object)
    self:Insert(object, position)
end


function Grid:RadiusSearch(from: Vector3, radius: number): ({ any }, { number })
    local foundObjects = {}
    local foundDistances = {}

    local radiusSquared = (radius * radius)
    local cellRadius = math.ceil(radius / self._cellSize)
    local fromCellX, fromCellZ = positionToCellPosition(from, self._cellSize)

    for positionX = fromCellX - cellRadius, fromCellX + cellRadius do
        for positionZ = fromCellZ - cellRadius, fromCellZ + cellRadius do
            local cellId = cellPositionToCellId(positionX, positionZ)

            if not self._cells[cellId] then
                continue
            end

            for object, position in pairs(self._cells[cellId]) do
                local distanceX = from.X - position.X
                local distanceZ = from.Z - position.Z
                local distanceSquared = (distanceX * distanceX) + (distanceZ * distanceZ)

                if distanceSquared <= radiusSquared then
                    table.insert(foundObjects, object)
                    table.insert(foundDistances, distanceSquared)
                end
            end
        end
    end

    return foundObjects, foundDistances
end

function Grid:RadiusSearchAmount(from: Vector3, radius: number, amount: number): ({ any }, { number })
    local objects, distances = self:RadiusSearch(from, radius)
    local sorted = {}
    
    for index, object in ipairs(objects) do
        local distanceSquared = distances[index]

        table.insert(sorted, {
            distanceSquared = distanceSquared,
            object = object
        })
    end

    table.sort(sorted, function(a, b)
        return a.distanceSquared < b.distanceSquared
    end)

    local foundObjects = {}
    local foundDistances = {}

    for i = 1, math.min(amount, #sorted) do
        local sortedEntry = sorted[i]
        foundObjects[i] = sortedEntry.object
        foundDistances[i] = sortedEntry.distanceSquared
    end

    return foundObjects, foundDistances
end

return Grid
1 Like
#2

I think this would be a great resource to have, you should add it under the community rescources. @LuxSkyWalkSchool

1 Like
#3

there are several bugs in the program though, and a few optimizations you could make. One bug is that you do both object[position] and position[object]. One optimization is the update function, maybye actually updating instead of destroying and then creating. @LuxSkyWalkSchool

1 Like
#4

Thank you for pointing out this bug. As for the optimizations, I’ll need to check if the Luau bytecode compiler is capable of inlining the function calls within the Update function, and if not possible manually inline it.