euluna raytracer test

raytrace.lua

!!strict

-- optimize for performance
!!cflags '-O3 -ffast-math -march=native -fno-plt -flto -fopenmp -rdynamic -g'
!!ldflags '-flto'

--------------------------------------------------------------------------------
-- SDL

-- import SDL headers
!!cdefine 'SDL_DISABLE_IMMINTRIN_H'
!!cinclude '<SDL2/SDL.h>'
!!linklib 'SDL2'

-- import SDL types
local SDL_Keysym = @record {
  scancode: cint,
  sym: int32,
  mod: uint16,
  unused: uint32
}
local SDL_KeyboardEvent = @record {
  type: uint32,
  timestamp: uint32,
  windowID: uint32,
  state: uint8,
  repeated: uint8,
  padding: uint16,
  keysym: SDL_Keysym
}
local SDL_Surface = @record{
  flags: uint32,
  format: pointer,
  w: cint,
  h: cint,
  pitch: cint,
  pixels: pointer
}
local SDL_Event = @record{type: uint32, padding: array<byte, 56>}
local SDL_Rect = @record{x: cint, y: cint, w: cint, h: cint}
local SDL_Window = @record{}*
local SDL_Renderer = @record{}*
local SDL_Texture = @record{}*

-- import SDL constants
local SDL_INIT_VIDEO: uint32 !cimport
local SDL_WINDOWPOS_UNDEFINED: cint !cimport
local SDL_WINDOW_OPENGL: uint32 !cimport
local SDL_QUIT: uint32 !cimport
local SDL_KEYDOWN: uint32 !cimport
local SDLK_UP: int32 !cimport
local SDLK_DOWN: int32 !cimport
local SDLK_LEFT: int32 !cimport
local SDLK_RIGHT: int32 !cimport
local SDLK_a: int32 !cimport
local SDLK_w: int32 !cimport
local SDLK_s: int32 !cimport
local SDLK_d: int32 !cimport
local SDLK_e: int32 !cimport
local SDLK_q: int32 !cimport
local SDL_PIXELFORMAT_ARGB8888: uint32 !cimport
local SDL_BLENDMODE_NONE: int32 !cimport
local SDL_RENDERER_ACCELERATED: uint32 !cimport
local SDL_RENDERER_PRESENTVSYNC: uint32 !cimport
local SDL_RENDERER_SOFTWARE: uint32 !cimport
local SDL_TEXTUREACCESS_STREAMING: cint !cimport

-- import SDL functions
local function SDL_Init(flags: uint32): int32 !cimport end
local function SDL_CreateWindow(title: cstring, x: cint, y: cint, w: cint, h: cint, flags: uint32): SDL_Window !cimport end
local function SDL_Quit() !cimport end
local function SDL_DestroyWindow(window: SDL_Window) !cimport end
local function SDL_PollEvent(event: SDL_Event*): int32 !cimport end
local function SDL_GetTicks(): uint32 !cimport end
local function SDL_Delay(ms: uint32) !cimport end
local function SDL_CreateRGBSurfaceWithFormatFrom(pixels: pointer, width: cint,height: cint, depth: cint, pitch: cint,format: uint32): SDL_Surface !cimport end
local function SDL_FreeSurface(surface: SDL_Surface) !cimport end
local function SDL_LockSurface(surface: SDL_Surface) !cimport end
local function SDL_GetWindowSurface(window: SDL_Window): SDL_Surface !cimport end
local function SDL_BlitSurface(src: SDL_Surface, srcrect: SDL_Rect*, dst: SDL_Surface, dstrect: SDL_Rect*): cint !cimport end
local function SDL_UpdateWindowSurface(window: SDL_Window) !cimport end
local function SDL_CreateRenderer(window: SDL_Window, index: cint, flags: uint32): SDL_Renderer !cimport end
local function SDL_DestroyRenderer(renderer: SDL_Renderer) !cimport end
local function SDL_RenderPresent(renderer: SDL_Renderer) !cimport end
local function SDL_RenderClear(renderer: SDL_Renderer) !cimport end
local function SDL_CreateTexture(renderer: SDL_Renderer, format: uint32, access: cint, w: cint, h: cint): SDL_Texture !cimport end
local function SDL_DestroyTexture(texture: SDL_Texture) !cimport end
local function SDL_RenderCopy(renderer: SDL_Renderer, texture: SDL_Texture, srcrect: SDL_Rect*, dstrect: SDL_Rect*): cint !cimport end
local function SDL_LockTexture(texture: SDL_Texture, rect: SDL_Rect*, pixels: pointer*, pitch: cint*): cint !cimport end
local function SDL_UnlockTexture(texture: SDL_Texture) !cimport end
local function SDL_SetRenderDrawBlendMode(renderer: SDL_Renderer, blendMode: int32): cint !cimport end
local function SDL_SetTextureBlendMode(texture: SDL_Texture, blendMode: int32): cint !cimport end
local function SDL_UpdateTexture(texture: SDL_Texture, rect: SDL_Rect*, pixels: pointer, pitch: cint): cint !cimport end
local function SDL_GetError(): cstring !cimport end

--------------------------------------------------------------------------------
-- C functions
local function rand(): int32 !cimport('rand') end
local function memcpy(dest: pointer, src: pointer, n: csize): pointer !cimport end

## if true then
local float = @float64
local function math_sqrt(x: float64): float64 !cimport('sqrt') end
local function math_sin(x: float64): float64 !cimport('sin') end
local function math_cos(x: float64): float64 !cimport('cos') end
local function math_pow(a: float64, b: float64): float64 !cimport('pow') end
local function math_tan(x: float64): float64 !cimport('tan') end
## else
local float = @float32
local function math_sqrt(x: float32): float32 !cimport('sqrtf') end
local function math_sin(x: float32): float32 !cimport('sinf') end
local function math_cos(x: float32): float32 !cimport('cosf') end
local function math_pow(a: float32, b: float32): float32 !cimport('powf') end
local function math_tan(x: float32): float32 !cimport('tan') end
## end

--------------------------------------------------------------------------------
-- vec3

local vec3 !aligned(16) = @record{x: float, y: float, z: float}
local function vec3_add(a: vec3, b: vec3): vec3 !inline
  return vec3{a.x+b.x, a.y+b.y, a.z+b.z}
end

local function vec3_mul(a: vec3, b: vec3): vec3 !inline
  return vec3{a.x*b.x, a.y*b.y, a.z*b.z}
end

local function vec3_addmul(a: vec3, b: vec3, factor: float): vec3 !inline
  return vec3{a.x+factor*b.x, a.y+factor*b.y, a.z+factor*b.z}
end

local function vec3_sub(a: vec3, b: vec3): vec3 !inline
  return vec3{a.x-b.x, a.y-b.y, a.z-b.z}
end

local function vec3_neg(a: vec3): vec3 !inline
  return vec3{-a.x, -a.y, -a.z}
end

local function vec3_dot(a: vec3, b: vec3): float !inline
  return a.x*b.x + a.y*b.y + a.z*b.z
end

local function vec3_cross(a: vec3, b: vec3): vec3 !inline
  return vec3{a.y*b.z - a.z*b.y, a.z*b.x - a.x*b.z, a.x*b.y - a.y*b.x}
end

local function vec3_sqrt(a: vec3): vec3 !inline
  return vec3{math_sqrt(a.x), math_sqrt(a.y), math_sqrt(a.z)}
end

local function vec3_smul(a: vec3, factor: float): vec3 !inline
  return vec3{a.x*factor, a.y*factor, a.z*factor}
end

local function vec3_sdiv(a: vec3, factor: float): vec3 !inline
  local k = 1 / factor
  return vec3{a.x*k, a.y*k, a.z*k}
end

local function vec3_squaredlength(v: vec3): float !inline
  return v.x*v.x + v.y*v.y + v.z*v.z
end

local function vec3_length(v: vec3): float !inline
  return math_sqrt(vec3_squaredlength(v))
end

local function vec3_unit(v: vec3): vec3 !inline
  return vec3_sdiv(v, vec3_length(v))
end

local function vec3_lerp(a: vec3, b: vec3, t: float): vec3 !inline
  return vec3_add(vec3_smul(a, 1-t), vec3_smul(b, t))
end


--------------------------------------------------------------------------------
-- mat4

local mat4 !aligned(32) = @float[4][4]

local function mat4_eye(): mat4 !inline
  return mat4{
    {1,0,0,0},
    {0,1,0,0},
    {0,0,1,0},
    {0,0,0,1}
  }
end

local function mat4_mul(A: mat4, B: mat4): mat4
  local C: mat4
  ## for i=0,3 do
    ## for j=0,3 do
      ## for k=0,3 do
        C[#[i]][#[j]] = C[#[i]][#[j]] + A[#[i]][#[k]] * B[#[k]][#[j]]
      ## end
    ## end
  ## end
  return C
end

local function mat4_mulvec3(A: mat4, v: vec3): vec3
  return vec3{
    A[0][0] * v.x + A[0][1] * v.y + A[0][2] * v.z,
    A[1][0] * v.x + A[1][1] * v.y + A[1][2] * v.z,
    A[2][0] * v.x + A[2][1] * v.y + A[2][2] * v.z
  }
end

local function vec3_rotX(v: vec3, theta: float): vec3
  local c = math_cos(theta)
  local s = math_sin(theta)
  local R = mat4{
    { 1, 0, 0, 0},
    { 0, c,-s, 0},
    { 0, s, c, 0},
    { 0, 0, 0, 1}
  }
  return mat4_mulvec3(R, v)
end

local function vec3_rotY(v: vec3, theta: float): vec3
  local c = math_cos(theta)
  local s = math_sin(theta)
  local R = mat4{
    { c, 0, s, 0},
    { 0, 1, 0, 0},
    {-s, 0, c, 0},
    { 0, 0, 0, 1}
  }
  return mat4_mulvec3(R, v)
end

local function vec3_rotZ(v: vec3, theta: float): vec3
  local c = math_cos(theta)
  local s = math_sin(theta)
  local R = mat4{
    { c,-s, 0, 0},
    { s, c, 0, 0},
    { 0, 0, 1, 0},
    { 0, 0, 0, 1}
  }
  return mat4_mulvec3(R, v)
end

--------------------------------------------------------------------------------

local MATERIALS = @enum {
  LAMBERTIAN = 0,
  METAL = 1,
  DIELECTRIC = 2
}
local vec4b = @record{x: byte, y: byte, z: byte, w: byte}
local ray = @record{origin: vec3, direction: vec3}
local material = @record{kind: MATERIALS, albedo: vec3, fuzz: float}
local hitrecord = @record{p: vec3, normal: vec3, t: float, mat: material*}
local sphere = @record{center: vec3, radius: float, mat: material}
local camera = @record{
  lower_left_corner: vec3,
  horizontal: vec3,
  vertical: vec3,
  origin: vec3,
  u: vec3,
  v: vec3,
  w: vec3
}

local const SCREEN_WIDTH: int32 = 960
local const SCREEN_HEIGHT: int32 = 540

local window, renderer, texture
local pixels: vec4b[SCREEN_WIDTH][SCREEN_HEIGHT] !aligned(16)
local const worldsize = 5
local hitablelist = @sphere[worldsize]
local world: hitablelist
world[0] = sphere{center={0,0,-1}, radius=0.5, mat={MATERIALS.LAMBERTIAN, {0.1, 0.2, 0.5}}}
world[1] = sphere{center={0,-100.5,-1}, radius=100, mat={MATERIALS.LAMBERTIAN, {0.8, 0.8, 0.0}}}
world[2] = sphere{center={1,0,-1}, radius=0.5, mat={MATERIALS.METAL, {0.8, 0.6, 0.2}, 0.0}}
world[3] = sphere{center={-1,0,-1}, radius=0.5, mat={MATERIALS.DIELECTRIC, fuzz=1.5}}
world[4] = sphere{center={-1,0,-1}, radius=-0.45, mat={MATERIALS.DIELECTRIC, fuzz=1.5}}


local function ray_pointat(r: ray, t: float): vec3 !inline
  return vec3_addmul(r.origin, r.direction, t)
end

local function camera_getray(self: camera*, u: float, v: float)
  return ray {
    origin = self.origin,
    direction = vec3_sub(vec3_add(self.lower_left_corner,
      vec3_add(vec3_smul(self.horizontal, u), vec3_smul(self.vertical, v))), self.origin)
  }
end

local function sphere_hit(self: sphere*, r: ray, tmin: float, tmax: float): boolean, hitrecord
  local oc = vec3_sub(r.origin, self.center)
  local a = vec3_dot(r.direction, r.direction)
  local b = vec3_dot(oc, r.direction)
  local c = vec3_dot(oc, oc) - self.radius*self.radius
  local discriminant = b*b - a*c
  local rec: hitrecord
  local t: float
  if discriminant > 0 then
    local droot = math_sqrt(discriminant)
    local ainv = 1 / a
    t = (-b - droot) * ainv
    if t < tmax and t > tmin then
      goto hitted
    end
    t = (-b + droot) * ainv
    if t < tmax and t > tmin then
      goto hitted
    end
  end
  do return false, rec end
::hitted::
  rec.t = t
  rec.p = ray_pointat(r, t)
  rec.normal = vec3_sdiv(vec3_sub(rec.p, self.center), self.radius)
  rec.mat = &self.mat
  return true, rec
end

local function hitablelist_hit(self: hitablelist*, r: ray, tmin: float, tmax: float): boolean, hitrecord
  local rec: hitrecord
  local tnear = tmax
  local hashit = false
  for i=0,<worldsize do
    local hitted, temprec = sphere_hit(&self[i], r, tmin, tnear)
    if hitted then
      hashit = true
      tnear = temprec.t
      rec = temprec
    end
  end
  return hashit, rec
end

local seed: int32 = 0
local function frand(): float
  seed = (214013*seed+2531011)
  return ((seed >> 16) & 32767) / @float(32768)
end

local function random_in_unit_sphere(): vec3
  local p: vec3
  repeat
    p = vec3_sub(vec3_smul(vec3{frand(), frand(), frand()}, 2), vec3{1,1,1})
  until vec3_squaredlength(p) < 1 and vec3_squaredlength(p) > 0.1
  return p
end

local function material_lambertian_scatter(self: material*, r: ray, rec: hitrecord): boolean, vec3, ray
  local target = vec3_add(vec3_add(rec.p, rec.normal), random_in_unit_sphere())
  local scattered = ray{rec.p, vec3_sub(target, rec.p)}
  return true, self.albedo, scattered
end

local function reflect(v: vec3, n: vec3): vec3
  return vec3_sub(v, vec3_smul(n, 2*vec3_dot(v,n)))
end

local function material_metal_scatter(self: material*, r: ray, rec: hitrecord): boolean, vec3, ray
  local reflected = reflect(vec3_unit(r.direction), rec.normal)
  if self.fuzz > 0 then
    reflected = vec3_add(reflected, vec3_smul(random_in_unit_sphere(), self.fuzz))
  end
  local scattered = ray{rec.p, reflected}
  local forward = vec3_dot(scattered.direction, rec.normal) > 0
  return forward, self.albedo, scattered
end

local function refract(v: vec3, n: vec3, ni: float): boolean, vec3
  local uv = vec3_unit(v)
  local dt = vec3_dot(uv, n)
  local discriminant = 1 - ni*ni*(1-dt*dt)
  if discriminant > 0 then
    local refracted = vec3_sub(
      vec3_smul(vec3_sub(uv, vec3_smul(n, dt)), ni),
      vec3_smul(n, math_sqrt(discriminant)))
    return true, refracted
  else
    return false, vec3{}
  end
end

local function schlick(cosine: float, refidx: float): float
  local r0 = (1-refidx) / (1+refidx)
  r0 = r0*r0
  return r0 + (1-r0)*math_pow(1-cosine, 5)
end

local function material_dielectric_scatter(self: material*, r: ray, rec: hitrecord): boolean, vec3, ray
  local reflected = reflect(r.direction, rec.normal)
  local outward_normal: vec3
  local ni: float, reflectprob: float
  local cosine = vec3_dot(r.direction, rec.normal) / vec3_length(r.direction)
  if vec3_dot(r.direction, rec.normal) > 0 then
    cosine = math_sqrt(1 - self.fuzz*self.fuzz*(1-cosine*cosine))
    outward_normal = vec3_neg(rec.normal)
    ni = self.fuzz
  else
    cosine = -cosine
    outward_normal = rec.normal
    ni = 1 / self.fuzz
  end
  local forward, refracted = refract(r.direction, outward_normal, ni)
  local scattered = ray{rec.p}
  if forward then
    reflectprob = schlick(cosine, self.fuzz)
  else
    reflectprob = 1
  end
  if frand() < reflectprob then
    scattered.direction = reflected
  else
    scattered.direction = refracted
  end
  return true, vec3{1,1,1}, scattered
end

local function material_scatter(self: material*, r: ray, rec: hitrecord): boolean, vec3, ray
  switch self.kind
  case MATERIALS.LAMBERTIAN then
    return material_lambertian_scatter(self, r, rec)
  case MATERIALS.METAL then
    return material_metal_scatter(self, r, rec)
  case MATERIALS.DIELECTRIC then
    return material_dielectric_scatter(self, r, rec)
  end
  return false, vec3{}, ray{}
end

local function raycast(r: ray, depth: integer): vec3
  local hitted, rec = hitablelist_hit(&world, r, 0.001, 1e32)
  if hitted then
    if depth > 16 then
      return vec3{0,0,0}
    end
    local forward, attenuation, scattered = material_scatter(rec.mat, r, rec)
    if forward then
      return vec3_mul(raycast(scattered, depth+1), attenuation)
    else
      return vec3{0,0,0}
    end
  end
  local unit_direction = vec3_unit(r.direction)
  local t = 0.5*(unit_direction.y + 1)
  return vec3_lerp(vec3{1,1,1}, vec3{0.5, 0.7, 1}, t)
end

local const math_pi: float = 3.141592653589793
local function camera_create(lookfrom: vec3, lookat: vec3, vup: vec3, vfov: float, aspect: float)
  local self = camera{}
  local theta = vfov*math_pi/180
  local half_height = math_tan(theta/2)
  local half_width = aspect * half_height
  self.w = vec3_unit(vec3_sub(lookfrom, lookat))
  self.u = vec3_unit(vec3_cross(vup, self.w))
  self.v = vec3_cross(self.w, self.u)
  self.lower_left_corner = vec3_sub(lookfrom, vec3_add(vec3_add(vec3_smul(self.u, half_width), vec3_smul(self.v, half_height)), self.w))
  self.horizontal = vec3_smul(self.u, 2*half_width)
  self.vertical = vec3_smul(self.v, 2*half_height)
  self.origin = lookfrom
  return self
end

local const SN = 2
local const SOFFS: float[5] = {-0.5,-0.33,0,0.33,0.5}

local lookfrom = vec3{0,0,0}
local lookat = vec3{0,0,-1}
local vup = vec3{0,1,0}
local fov = 59
local aspect = SCREEN_WIDTH/@float(SCREEN_HEIGHT)
local cam = camera_create(lookfrom, lookat, vup, fov, aspect)

local function camera_update()
  cam = camera_create(lookfrom, lookat, vup, fov, aspect)
end

local function draw_rays()
  !!cemit '#pragma omp parallel for schedule(dynamic)'
  for y=0,<SCREEN_HEIGHT do
    for x=0,<SCREEN_WIDTH do
      local col: vec3
      for sy=0,2*SN do
        local v = (y + SOFFS[sy]) / @float(SCREEN_HEIGHT)
        for sx=0,2*SN do
          local u = (x + SOFFS[sx]) / @float(SCREEN_WIDTH)
          col = vec3_add(col, raycast(camera_getray(&cam, u, v), 0))
        end
      end
      col = vec3_sdiv(col, (SN*2+1)*(SN*2+1))
      col = vec3_sqrt(col) -- gamma correction
      local colb = vec4b{@uint8(col.z*255), @uint8(col.y*255), @uint8(col.x*255), 255}
      pixels[SCREEN_HEIGHT - y - 1][x] = colb
    end
  end
end

local function upload_pixels()
  SDL_UpdateTexture(texture, nilptr, &pixels[0][0], @cint(SCREEN_WIDTH*4))
  SDL_RenderCopy(renderer, texture, nilptr, nilptr)
  SDL_RenderPresent(renderer)
end

local function draw()
  draw_rays()
  upload_pixels()
end

local function camera_rotate(x: float, y: float, z: float)
  local dir = vec3_sub(lookat, lookfrom)
  dir = vec3_rotX(dir, x)
  dir = vec3_rotY(dir, y)
  dir = vec3_rotZ(dir, z)
  lookat = vec3_add(lookfrom, dir)
  camera_update()
end

local function camera_translate(x: float, y: float, z: float)
  local trans = vec3{0,0,0}
  trans = vec3_add(trans, vec3_smul(cam.u, x))
  trans = vec3_add(trans, vec3_smul(cam.w,-z))
  trans = vec3_add(trans, vec3_smul(cam.v, y))
  lookfrom = vec3_add(lookfrom, trans)
  lookat = vec3_add(lookat, trans)
  camera_update()
end

local function poll_events()
  local event: SDL_Event
  while SDL_PollEvent(&event) ~= 0 do
    switch event.type
    case SDL_QUIT then
      return false
    case SDL_KEYDOWN then
      local kevent = @SDL_KeyboardEvent*(&event)
      switch kevent.keysym.sym
      case SDLK_UP then
        camera_rotate(0.1, 0, 0)
      case SDLK_DOWN then
        camera_rotate(-0.1, 0, 0)
      case SDLK_RIGHT then
        camera_rotate(0, -0.1, 0)
      case SDLK_LEFT then
        camera_rotate(0, 0.1, 0)
      case SDLK_w then
        camera_translate(0, 0, 0.1)
      case SDLK_s then
        camera_translate(0, 0, -0.1)
      case SDLK_a then
        camera_translate(-0.1, 0, 0)
      case SDLK_d then
        camera_translate(0.1, 0, 0)
      case SDLK_e then
        camera_translate(0, 0.1, 0)
      case SDLK_q then
        camera_translate(0, -0.1, 0)
      end
    end
  end
  return true
end

local function go()
  -- init sdl
  SDL_Init(SDL_INIT_VIDEO)

  window = SDL_CreateWindow("An SDL2 Window",
    SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED,
    SCREEN_WIDTH, SCREEN_HEIGHT, SDL_WINDOW_OPENGL)
  assert(window, "Could not create window")

  renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED)
  assert(renderer, "Could not create renderer")
  SDL_SetRenderDrawBlendMode(renderer, SDL_BLENDMODE_NONE)

  texture = SDL_CreateTexture(renderer, SDL_PIXELFORMAT_ARGB8888, SDL_TEXTUREACCESS_STREAMING, SCREEN_WIDTH, SCREEN_HEIGHT)
  SDL_SetTextureBlendMode(texture, SDL_BLENDMODE_NONE)

  -- draw
  local lastticks = SDL_GetTicks()
  local fps = 0
  repeat
    local ticks = SDL_GetTicks()
    if ticks - lastticks >= 1000 then
      print('FPS', fps)
      lastticks = ticks
      fps = 0
    end

    local quit = not poll_events()
    draw()
    fps = fps + 1
  until quit

  -- cleanup and finish
  SDL_DestroyTexture(texture)
  SDL_DestroyRenderer(renderer)
  SDL_DestroyWindow(window)
  SDL_Quit()
end

go()

--TODO: multidimensional arrays shortcut swap?
--TODO: record methods
--TODO: record field aliases
--TODO: record operators
--TODO: record overloading
--TODO: math library
--TODO: evaluate const numeric expressions
--TODO: automatic ref/deref
--TODO: optional type
--TODO: remove !! pragmas and use preprocessor instead