Orbital: support multiple static bodies

lean/ECS/System.lean

@@ -16,7 +16,7 @@ def get
     let s ← Has.getStore c
     comp.constraint.mp <$> e.explGet s ety
 
--- TODO: I want to call this `set` but it condflicts with a Prelude function
+-- TODO: I want to call this `set` but it conflicts with a Prelude function
 /-- Writes a component to the given entity. Will overwrite existing components --/
 def set'
   {c s t : Type}
@@ -40,7 +40,7 @@ def exists?
   [e : @ExplGet s t _]
   (ety : Entity) : System w Bool := do
   let s ← Has.getStore c
-  e.explExists  s ety
+  e.explExists s ety
 
 /-- Destroys component c for the given enitty --/
 def destroy
@@ -75,6 +75,19 @@ def modify'
       let x ← getX.explGet sx ety
       setY.explSet sy ety (compY.constraint.symm.mp (f (compX.constraint.mp x)))
 
+/-- Returns an array of all components with type cx --/
+def members
+  [FamilyDef StorageFam cx sx]
+  [FamilyDef ElemFam sx tx]
+  [compX : @Component cx sx tx _ _]
+  [@Has w cx sx _]
+  [getX : @ExplGet sx tx _]
+  [mX : ExplMembers sx] : System w (Array cx) := do
+    let stx ← Has.getStore cx
+    let sl ← mX.explMembers stx
+    let res : Array tx ← sl.mapM (getX.explGet stx) |> monadLift
+    return (res.map compX.constraint.mp)
+
 /-- Maps a function over all entities with a cx component and writes their cy --/
 def cmap
   [FamilyDef StorageFam cx sx]

lean/Examples/Orbital.lean

@@ -24,6 +24,8 @@ def mkStaticBody (position : Vector2) : System World Unit :=
 def init : System World Unit := do
 
   mkStaticBody origin
+  mkStaticBody (origin.add ⟨3, 3⟩)
+  mkStaticBody (origin.add ⟨-1, 3⟩)
   mkOrbitingBody (selected := true) initPos initVel
   mkOrbitingBody ⟨-4, 0⟩ (initVel.mul (-1))
   mkOrbitingBody ⟨-0.8, -0.6⟩ (initVel.mul (-0.95))
@@ -34,15 +36,24 @@ def init : System World Unit := do
 
 def terminate : System World Unit := closeWindow
 
-def updateOrbitingBody (dt : Float) : Position × Velocity × OrbitPath → Position × Velocity × OrbitPath
-  | (Position.mk p, Velocity.mk v, OrbitPath.mk o) =>
-    let pMag := p.length
-    let a := p |>.mul (-1 / pMag^3)
-    let vNew := v.add (a.mul dt)
-    let pNew := p.add (vNew.mul dt)
+def updateWithStatic (dt : Float) (staticBodies : Array Position) : Position × Velocity × OrbitPath → Position × Velocity × OrbitPath
+  | ⟨⟨po⟩, ⟨v⟩, ⟨o⟩⟩ => Id.run do
+    -- compute the new velocity from contributions from each static body
+    let mut vNew := v
+    for ⟨ps⟩ in staticBodies do
+      -- p is the vector pointing from the oribiting body (po) to the static body (ps)
+      let p := ps.sub po
+      let pMag := p.length
+      let a := p |>.mul (1 / pMag^3)
+      vNew := vNew.add (a.mul dt)
+    let pNew := po.add (vNew.mul dt)
     let oNew := o.push pNew
     (⟨pNew⟩, ⟨vNew⟩, ⟨oNew⟩)
 
+def updateOrbitingBody (dt : Float) : System World Unit := do
+  let static : Array (Position × Not Velocity) ← members
+  cmap (updateWithStatic dt (static.map Prod.fst))
+
 def changeSelectedVelocity (dv : Float) : Velocity × Selectable → Velocity
  | (Velocity.mk v, Selectable.mk true) => ⟨v.add <| v.mul (dv / v.length)⟩
  | (Velocity.mk v, Selectable.mk false) => ⟨v⟩
@@ -66,14 +77,14 @@ def update : System World Unit := do
   if (← isKeyDown Key.left) then cmap (changePerpVelocity (-0.01))
   if (← isKeyDown Key.space) then cmap resetOrbitPath
   if (← isKeyDown Key.r) then cmap resetSelected
-  cmap (updateOrbitingBody (← getFrameTime))
+  updateOrbitingBody (← getFrameTime)
 
 /-- Convert a Position to a point on the screen --/
 def toScreen (v : Vector2) : Vector2 :=
   v.mul 100 |>.add center
 
 def renderStaticBody (p : Position) : System World Unit :=
-    drawCircleV (p.val.add center) 30 Color.red
+    drawCircleV (toScreen p.val) 30 Color.red
 
 def renderOrbitingBody (p : Position) (s : Selectable) : System World Unit :=
   drawCircleV (toScreen p.val) 10 (if s.selected then Color.green else Color.blue)