Imperative style update syntax for FSharp.Data JsonValue using lenses

0_Usage.fs

let jsonStr = """
{
    "itemSize": {
        "width": "hello"
    },
    "tags": ["foo","bar"]
}
"""

let json = JsonValue.Parse jsonStr

let width = 
    JsonLens.recordPL >=> 
    JsonLens.recordFieldPL "itemSize" >=> 
    JsonLens.recordPL >=> 
    JsonLens.recordFieldPL "width" >=>
    JsonLens.stringPL

let firstTag =
    JsonLens.recordPL >=>
    JsonLens.recordFieldPL "tags" >=>
    JsonLens.arrayPL >=>
    JsonLens.arrayItemPL 0

let json' = json |> State.exec (state {
    do! width := "world"
    do! firstTag := JsonValue.String "baz"
})

printfn "%A" json'

1_State.fs

/// The state monad - a transition from a state, to a value and a new state.
type State<'s, 'a> = State of ('s -> 'a * 's)

/// The state monad.
module State =   

    let inline state (f:'s -> 'a * 's) : State<'s, 'a> = State f

    let inline run s (state:State<'s, 'a>) = let (State(run)) = state in run s

    let inline eval (st:State<'s, 'a>) s = run s st |> fst

    let inline exec (st:State<'s, 'a>) s = run s st |> snd

    let inline delay (f:unit -> State<'s, 'a>) : State<'s, 'a> =
        state <| fun s -> f() |> run s

    let unit a : State<'s, 'a> = state <| fun s -> (a,s)

    let get<'s> : State<'s, 's> = state <| fun s -> (s,s)

    let set s : State<'s, unit> = state <| fun _ -> ((),s)

    let map f (st:State<'s, 'a>) : State<'s, 'b> = 
        state <| fun s -> 
            let (a,s1) = run s st
            (f a, s1)
    
    let bind f (st:State<'s, 'a>) : State<'s, 'b> =
        state <| fun s ->
            let (a,s1) = run s st
            run s1 (f a)
    
    let map2 (s1:State<'s, 'a>) (s2:State<'s, 'b>) (f:'a -> 'b -> 'c) : State<'s, 'c> =
        bind (fun a -> map (fun b -> f a b) s2) s1


    let traverseSeq (f:'a -> State<'State, 'b>) (ss:seq<'a>) : State<'State, 'b seq> =
        state <| fun s ->            
            let s = ref s
            let ss = 
                ss 
                |> Seq.map (fun a ->                     
                    let a,s' = f a |> run !s
                    s := s'
                    a
                )
                |> Seq.toArray
            (ss |> Seq.ofArray,!s)

    let sequenceSeq (ss:State<'s, 'a> seq) : State<'s, 'a seq> =
        traverseSeq id ss            


    let concatSeq (ss:seq<State<'s, 'a>>) (M:Monoid<'a>) : State<'s, 'a> = 
        ss |> sequenceSeq |> map (Seq.fold M.op M.unit)

    
    /// Creates a state transition which evaluates the array of state transitions and returns the array of results.
    let traverseArray (f:'a -> State<'S, 'b>) (xs:'a[]) : State<'S, 'b[]> =
        state <| fun s ->
            let s = ref s
            let rs =
                xs
                |> Array.map (fun a ->
                    let a,s' = f a |> run !s
                    s := s'
                    a
                )
            rs,!s
    
    let sequenceArray (xs:State<'S, 'a>[]) : State<'S, 'a[]> =
        traverseArray id xs


    
    /// Creates a state transition which evaluates the list of state transitions and returns the list of results.
    let traverseList (f:'a -> State<'S, 'b>) (xs:list<'a>) : State<'S, list<'b>> =
        state <| fun s ->
            let s = ref s
            let rs =
                xs
                |> List.map (fun a ->
                    let a,s' = f a |> run !s
                    s := s'
                    a
                )
            rs,!s

    let sequenceList (ss:State<'s, 'a> list) : State<'s, 'a list> =
        traverseList id ss


    type StateBuilder() =
        member __.Bind(s, f) = bind f s
        member __.Return(value) = unit value
        member __.Yield(value) = unit value
        member __.ReturnFrom(value) = value
        member __.Zero() = unit()
        member __.Delay(f) = delay f
        member __.Combine(s1:State<'S,unit>, s2:State<'S,'a>) = map2 s1 s2 (fun _ s -> s)
        member __.For(xs:seq<'a>, f:'a -> State<'S, 'a>) = xs |> Seq.map f


[<AutoOpen>]
module StateBuilder =

    /// State monad workflow builder.
    let state = new State.StateBuilder()






/// CoState comonad.
type CoState<'s, 'a> = 's * ('s -> 'a)
 
module CoState =
    
    let inline create s set : CoState<'s, 'a> = (s,set)
 
    let pos (st:CoState<'s, 'a>) : 's = let (s,_) = st in s
 
    let peek s (st:CoState<'s, 'a>) : 'a = 
        let (_,set) = st in
        set s

    let seek s (st:CoState<'s, 'a>) : CoState<'s, 'a> =
        let (_,set) = st in
        (s, set)

    let set (st:CoState<'s, 'a>) : 's -> 'a = let (_,set) = st in set
 
    let set_ (s:'s) (st:CoState<'s, 'a>) : 'a =
        set st s

    let map f (st:CoState<'s, 'a>) : CoState<'s, 'b> =
        let (s,set) = st in
        (s, f << set)
 
    let cobind (f:CoState<'s, 'a> -> 'b) (st:CoState<'s, 'a>) : CoState<'s, 'b> =
        let (s,set) = st in
        (s, fun b -> f (b,set))
        
    let cojoin (st:CoState<'s, 'a>) : CoState<'s, CoState<'s, 'a>> =
        cobind (Operators.id) st

    let experimentList (f:'s -> list<'s>) (st:CoState<'s, 'a>) : list<'a> =
        let (s,set) = st in
        f s |> List.map set

    let experimentArray (f:'s -> 's[]) (st:CoState<'s, 'a>) : 'a[] =
        let (s,set) = st in
        f s |> Array.map set

    let experimentOption (f:'s -> 's option) (st:CoState<'s, 'a>) : 'a option =
        let (s,set) = st in
        f s |> Option.map set
        

2_Lens.fs

// A lens is an abstraction of a property with get/set allowings for composition.
type Lens<'a, 'b> = {
    get: 'a -> 'b
    set: 'b -> 'a -> 'a 
}
with member l.update f a = let b = l.get a |> f in l.set b a
    

module Lens =    

    let inline create get set = { Lens.get = get ; set = set }        

    let inline get a (l:Lens<'a, 'b>) = l.get a
    
    let inline set b a (l:Lens<'a, 'b>) = l.set b a          
    
    let toCoState (l:Lens<'a, 'b>) : 'a -> CoState<'b, 'a> =
        fun a -> CoState.create (l.get a) (fun b -> l.set b a)

    let ofCoState (f:'a -> CoState<'b, 'a>) : Lens<'a, 'b> =
        create (f >> CoState.pos) (fun b a -> CoState.set (f a) b)

    /// Creates an isomorphism for a lens given f and g.
    let inline iso f g (l:Lens<'a, 'b>) : Lens<'a, 'c> = { get = l.get >> f ; set = g >> l.set }    
    
    /// Composes two lenses into one.
    let inline compose (l1:Lens<'a, 'b>) (l2:Lens<'b, 'c>) : Lens<'a, 'c> =
        { get = l1.get >> l2.get
          set = l2.set >> l1.update }



module LensOps =
 
    let (>=>) (f:Lens<'a, 'b>) (g:Lens<'b, 'c>) : Lens<'a, 'c> = Lens.compose f g

    let (:=) (l:Lens<'a, 'b>) b : State<'a, unit> = 
        State.state <| fun a -> ((), Lens.set b a l)
 
    let inline (+=) (l:Lens<'a, 'b>) b : State<'a, unit> = 
        State.state <| fun a ->
            let w = (Lens.get a l) + b
            ((), Lens.set w a l)



module LensLaws =
    
    let lensLaw1 (l:Lens<'a, 'b>) a b = l.get (l.set b a) = b      
    


/// Partial lens.
type PLens<'a, 'b> = 'a -> CoState<'b, 'a> option

/// Partial lens.
module PLens =
    
    let ofLens (l:Lens<'a, 'b>) : PLens<'a, 'b> = fun a -> Some (Lens.toCoState l a)

    let get a (pl:PLens<'a, 'b>) : 'b option =
        pl a |> Option.map CoState.pos

    let set b a (pl:PLens<'a, 'b>) : 'a option =
        pl a |> Option.map (CoState.set_ b)

    let create (get:'a -> 'b option) (set:'b -> 'a -> 'a option) : PLens<'a, 'b> =
        fun a -> 
            match get a with
            | Some b -> CoState.create b (fun b -> set b a |> Option.getValueOr a) |> Some
            | None -> None

    let left<'a, 'b> : PLens<Choice<'a, 'b>, 'a> =
        function
        | Choice1Of2 a -> CoState.create a Choice1Of2 |> Some
        | Choice2Of2 _ -> None

    let right<'a, 'b> : PLens<Choice<'a, 'b>, 'b> =
        function
        | Choice1Of2 _ -> None
        | Choice2Of2 b -> CoState.create b Choice2Of2 |> Some

    let listHead<'a> : PLens<list<'a>, 'a> =
        function
        | h::t -> CoState.create h (fun a -> a::t) |> Some
        | [] -> None

    let some<'a> : PLens<option<'a>, 'a> =
        function
        | Some a -> CoState.create a (Some) |> Some
        | None -> None

    let compose (f:PLens<'b, 'c>) (g:PLens<'a, 'b>) : PLens<'a, 'c> =
        fun a -> 
            g a
            |> Option.bind (fun ba ->
                f (CoState.pos ba)
                |> Option.map (fun cb ->
                    CoState.create (CoState.pos cb) (fun c -> CoState.set ba (CoState.set cb c))
                )
            )

    let arrayItem (i:int) : PLens<'a[], 'a> =
        create (Array.tryNth i) (fun a arr -> arr.[i] <- a ; Some arr)

            
module PLensOps =
    
    let (>=>) (f:PLens<'a, 'b>) (g:PLens<'b, 'c>) : PLens<'a, 'c> = PLens.compose g f

    let (:=) (l:PLens<'a, 'b>) b : State<'a, unit> = 
        State.state <| fun a -> ((), PLens.set b a l |> Option.getValueOr a)
 
    let inline (+=) (l:PLens<'a, 'b>) (b:'b) : State<'a, unit> = 
        State.state <| fun a ->
            let w = (PLens.get a l) |> Option.map ((+) b)
            match w with
            | Some w -> ((), PLens.set w a l |> Option.getValueOr a)
            | None -> ((), a)

3_JsonLens.fs

module internal JValue =
    
    let tryString = function
        | JsonValue.String str -> Some str
        | _ -> None      

    let tryRecord = function
        | JsonValue.Record props -> Some props
        | _ -> None

    let tryArray = function
        | JsonValue.Array els -> Some els
        | _ -> None

    let private fieldIndex (field:string) (o:(string * JsonValue)[]) : int option =
        o |> Array.tryFindIndex (fun (f,_) -> f = field)

    let getField (o:(string * JsonValue)[]) (field:string) : JsonValue option =
        o
        |> Array.tryFind (fun (f,_) -> f = field)
        |> Option.map snd 

    let removeField (o:(string * JsonValue)[]) (field:string) : (string * JsonValue)[] =
        match o |> fieldIndex field with
        | Some i -> Array.removeAt o i
        | None -> o            

    let setField (o:(string * JsonValue)[]) (field:string) (value:JsonValue) : (string * JsonValue)[] =
        match o |> fieldIndex field with
        | Some i ->
            o.[i] <- field,value
            o
        | None ->
            Array.snoc o (field,value)
            


/// JsonValue lens combinators.
module JsonLens =
  
    let stringPL : PLens<JsonValue, string> =
        JValue.tryString >> Option.map (fun str -> CoState.create str JsonValue.String)

    let recordPL : PLens<JsonValue, (string * JsonValue)[]> =
        JValue.tryRecord >> Option.map (fun props -> CoState.create props JsonValue.Record)

    let recordFieldL (field:string) : Lens<(string * JsonValue)[], JsonValue option> =
        Lens.create 
            (fun o -> JValue.getField o field)
            (fun f o -> match f with Some f -> JValue.setField o field f | None -> JValue.removeField o field)

    let recordFieldPL (field:string) : PLens<(string * JsonValue)[], JsonValue> =                
        recordFieldL field 
        |> PLens.ofLens
        |> PLens.compose PLens.some

    let arrayPL : PLens<JsonValue, JsonValue[]> =
        JValue.tryArray >> Option.map (fun arr -> CoState.create arr JsonValue.Array)
            

    let arrayItemPL (i:int) : PLens<JsonValue[], JsonValue> = PLens.arrayItem i