ECS in Nim

entity.nim

#[
  Copyright (c) 2019 RSDuck/Kemal Afzal

  Permission is hereby granted, free of charge, to any person obtaining a copy of this software 
  and associated documentation files (the "Software"), to deal in the Software without restriction, 
  including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, 
  and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, 
  subject to the following conditions:

  The above copyright notice and this permission notice shall be included in all copies or 
  substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 
  INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE 
  AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, 
  DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
]#

import
    bitops, macros, tables, macrocache, strformat, typetraits

{.experimental: "forLoopMacros".}

type
    DynBitSet = distinct seq[uint64]

    ComponentsRoot = ref object of RootObj
        used: DynBitSet
        mapping: seq[uint16]
        freeList: seq[uint16]
    Components[T] = ref object of ComponentsRoot
        data: seq[T] # component -> actual component data

    World = object
        usedEntities: DynBitSet
        generation: seq[uint8]

        components: seq[ComponentsRoot]

    Entity = distinct uint32

proc contains(self: DynBitSet, idx: uint32): bool =
    if int(idx) >= seq[uint64](self).len * 64:
        return false
    (seq[uint64](self)[idx div 64] and (1'u64 shl (idx mod 64))) != 0
proc `[]`(self: DynBitSet, idx: uint32): bool =
    idx in self
proc `[]=`(self: var DynBitSet, idx: uint32, val: bool) =
    let
        part = idx div 64
        subpart = idx mod 64
    if int(part) >= seq[uint64](self).len:
        seq[uint64](self).setLen(part + 1)
    seq[uint64](self)[part] = (seq[uint64](self)[part] and
        not(1'u64 shl subpart)) or (uint64(val) shl subpart)

proc findFirstFree(self: DynBitSet): uint32 =
    for i, part in pairs(seq[uint64](self)):
        if part != not(0'u64):
            return uint32(firstSetBit(not part)) - 1 + uint32(i) * 64
    uint32(seq[uint64](self).len * 64)

proc union*(a, b: DynBitSet): DynBitSet =
    result = DynBitSet(newSeq[uint64](min(seq[uint64](a).len, seq[uint64](b).len)))
    for i in 0..<seq[uint64](result).len:
        seq[uint64](result)[i] = seq[uint64](a)[i] and seq[uint64](b)[i]

iterator items*(self: DynBitSet): uint32 =
    var i = 0
    while i < seq[uint64](self).len:
        var part = seq[uint64](self)[i]
        while part != 0:
            let nextItem = firstSetBit(part) - 1
            part = part and not(1'u64 shl nextItem)
            yield uint32(i * 64 + nextItem)

        inc i

proc idx(self: Entity): uint32 =
    uint32(self) and 0xffff_ff
proc generation(self: Entity): uint8 =
    uint8(uint32(self) shr 24)
proc initEntity(id: uint32, gen: uint8): Entity =
    assert id <= ((1 shl 24) - 1)
    assert gen <= ((1 shl 8) - 1)

    Entity(id or (gen shl 24))
proc `$`*(self: Entity): string =
    fmt"(idx: {self.idx}, generation: {self.generation})"

var componentTypes {.compiletime.} = initTable[string, int]()

macro typeIdx*(typ: typedesc): int =
    let name = typ.getTypeInst().repr()

    result = newIntLitNode componentTypes.getOrDefault(name, componentTypes.len)
    if result.intVal == componentTypes.len:
        componentTypes[name] = componentTypes.len

proc newEntity*(self: var World): Entity =
    let idx = self.usedEntities.findFirstFree()
    if self.generation.len <= int idx:
        self.generation.setLen(idx + 1)
        for components in self.components:
            components.mapping.setLen(idx + 1)

    self.usedEntities[idx] = true
    initEntity(idx, self.generation[idx])

proc isValid*(self: World, entity: Entity): bool =
    entity.idx in self.usedEntities and entity.generation == self.generation[entity.idx]

proc freeEntity*(self: var World, entity: Entity) =
    if self.isValid(entity):
        self.usedEntities[entity.idx] = false
        self.generation[entity.idx] = uint8((int(self.generation[entity.idx]) + 1) mod 256)

        for components in self.components:
            if components.used[entity.idx]:
                components.freeList.add(components.mapping[entity.idx])
                components.used[entity.idx] = false
    else:
        raise newException(KeyError, fmt"Invalid entity {entity}")

proc getComponents[T](self: var World): Components[T] =
    let idx = typeIdx(T)
    if self.components.len <= idx:
        self.components.setLen(idx + 1)
    if self.components[idx] == nil:
        self.components[idx] = new(Components[T])
        self.components[idx].mapping.setLen(self.generation.len)

    cast[Components[T]](self.components[idx])

proc hasComponent*[T](self: var World, entity: Entity): bool =
    if self.isValid(entity):
        let components = self.getComponents[:T]()
        return components.used[entity.idx]
    raise newException(KeyError, fmt"Invalid entity {entity}")

proc assignComponent*[T](self: var World, entity: Entity, component: T) =
    if self.isValid(entity):
        let components = self.getComponents[:T]()

        if not components.used[entity.idx]:
            if components.freeList.len == 0:
                let newLen = (components.data.len * 3) div 2 + 4
                for i in components.data.len..<newLen:
                    components.freeList.add(uint16 i)
                components.data.setLen(newLen)

            components.mapping[entity.idx] = components.freeList.pop()
            components.used[entity.idx] = true

        components.data[components.mapping[entity.idx]] = component
    else:
        raise newException(KeyError, fmt"Invalid entity {entity}")

proc getComponent*[T](self: var World, entity: Entity): T =
    if self.isValid(entity):
        var components = self.getComponents[:T]()

        if components.used[entity.idx]:
            return components.data[components.mapping[entity.idx]]
        else:
            raise newException(ValueError, fmt"Entity has no component of type {name(T)} attached")
    raise newException(KeyError, fmt"Invalid entity {entity}")

proc removeComponent*[T](self: var World, entity: Entity) =
    if self.isValid(entity):
        let components = self.getComponents[:T]()
        if components.used[entity.idx]:
            components.freeList.add(components.mapping[entity.idx])
            components.used[entity.idx] = false
    else:
        raise newException(KeyError, fmt"Invalid entity {entity}")

macro entitiesWith*(x: ForLoopStmt): untyped =
    let
        body = x[^1]
        self = x[^2][1]
        names = x[0..^3]
        types = x[^2][2..^1]

        typeIndices = nskConst.genSym("typeIndices")
        index = nskForVar.genSym("i")

        indices = nnkBracket.newTree()

        accessors = nnkVarSection.newTree()

        preparations = newStmtList()

    if names.len != types.len:
        error(fmt"Mismatch between named parameters({names.repr}) and desired types{types.repr}")

    for i in 0..<types.len:
        let typ = types[i]

        if typ != ident"Entity":
            preparations.add(quote do: discard `self`.getComponents[:`typ`]())
            indices.add(nnkCall.newTree(bindSym"typeIdx", typ))

            if names[i] != ident"_":
                accessors.add(nnkIdentDefs.newTree(names[i], newEmptyNode(), quote do:
                    addr cast[Components[`typ`]]
                        (`self`.components[`typeIndices`[`i`]]).data[`self`.components[`typeIndices`[`i`]].mapping[`index`]]))
        else:
            indices.add(newIntLitNode(0))

            accessors.add(nnkIdentDefs.newTree(names[i], newEmptyNode(), quote do:
                initEntity(`index`, `self`.generation[`index`])))

    quote do:
        block:
            `preparations`

            const `typeIndices` = `indices`
            var locations = `self`.components[`typeIndices`[0]].used
            for i in 0..<`self`.components.len:
                locations = union(locations, `self`.components[i].used)

            for `index` in locations:
                `accessors`
                `body`

when isMainModule:
    import random, intsets

    var
        world: World
        entities = newSeq[Entity]()

        values = initIntSet()
    for i in 0..<100:
        let e = world.newEntity()
        if i mod 2 == 0:
            world.assignComponent(e, i)
            values.incl i
        world.assignComponent(e, (20, 30))
            #echo world.getComponent[:int](e)
        #world.removeComponent[:int](e)
        entities.add(e)

    var i = 0
    for e, ci, _ in entitiesWith(world, Entity, int, (int, int)):
        doAssert ci[] in values
        values.excl ci[]
        world.assignComponent(e, ("miau", 42))
        inc i
    doAssert values.len == 0

    for _ in entitiesWith(world, (string, int)):
        dec i
    doAssert i == 0

    shuffle(entities)
    for e in entities:
        if world.hasComponent[:(string, int)](e):
            inc i
        world.freeEntity(e)
        if i == 50:
            break

    for _ in entitiesWith(world, (string, int)):
        dec i

    doAssert i == 50