ThreeVecs_State_help.jl

mutable struct ThreeVec <: AbstractVector{Float64}
    x::Float64
    y::Float64
    z::Float64
end

Base.size(V::ThreeVec) = (3,)
Base.IndexStyle(::Type{<:ThreeVec}) = IndexLinear()
using Match
Base.getindex(V::ThreeVec, i::Int) = @match i begin
    1 => V.x
    2 => V.y 
    3 => V.z
    _ => error("Index $is out of bounds")
end

function Base.setindex!(V::ThreeVec, val::Float64, i::Int) 
    if i==1
        V.x = val
    elseif i==2
        V.y = val
    elseif i==3
        V.z = val
    else
        error("Index $i is out of bounds")
    end
    return V
end

# State vector
mutable struct State <: AbstractMatrix{Float64}
    r::ThreeVec
    v::ThreeVec
end

Base.size(S::State) = (2,3)
Base.IndexStyle(::Type{<:State}) = IndexCartesian()
Base.getindex(S::State, I) = begin
    if I[1] == 1
        S.r[I[2]]
    elseif I[1] == 2
        S.v[I[2]]
    else
        error("Index $I out of bounds")
    end
end

Base.setindex!(S::State, val::Float64, I) = begin
    if I[1] == 1
        S.r[I[2]] = val
    elseif I[1] == 2
        S.v[I[2]] = val
    else
        error("Index $I out of bounds")
    end
    return S
end


using Test
@testset "Testing my code" begin
    v₁ = ThreeVec(1.0, 2.0, 3.0)
    v₂ = ThreeVec(4.0, 5.0, 6.0)
    s = State(v₁, v₂)
    @testset "Printing" begin
        println("v₁ = ", v₁)
        println("v₂ = ", v₂)
        println("s = ", s)
    end
    @testset "Accessing via indices" begin
        s[1,1] = 7.0
        @test s[1,1] == s.r.x
        @test s[1,2] == s.r.y
        @test s[1,3] == s.r.z
        @test s[2,1] == s.v.x
        @test s[2,2] == s.v.y
        @test s[2,3] == s.v.z
    end
    @testset "Multiplying by a number" begin
        a = 1.0
        s₂ = a * s
        @test s₂[2,1] == a * s[2,1]
    end
    @testset "Adding states" begin
        a = 1.0
        s₂ = a * s
        s₃ = s + s₂
        @test s₃[2,1] == s[2,1] + s₂[2,1]
    end
end
println("All tests finished.")