lenses.ts

class Lens<RootObject extends object, ChangableObject> {
  getter: (x: RootObject) => ChangableObject;
  setter: (x: RootObject, newValue: ChangableObject) => RootObject;

  static simple<RootObject extends object, T extends keyof RootObject>(
    name: T
  ) {
    return new Lens<RootObject, RootObject[T]>(
      x => x[name],
      (x, v) => ({ ...x, [name]: v })
    );
  }

  constructor(
    getter: Lens<RootObject, ChangableObject>["getter"],
    setter: Lens<RootObject, ChangableObject>["setter"]
  ) {
    this.getter = getter;
    this.setter = setter;
  }

  get(x: RootObject): ChangableObject {
    return this.getter(x);
  }

  set(x: RootObject, value: ChangableObject): RootObject {
    return this.setter(x, value);
  }

  over(
    x: RootObject,
    fn: (value: ChangableObject) => ChangableObject
  ): RootObject {
    return this.setter(x, fn(this.getter(x)));
  }

  compose<ChangableObject2>(
    l2: Lens<
      ChangableObject extends object ? ChangableObject : never,
      ChangableObject2
    >
  ): Lens<RootObject, ChangableObject2> {
    return new Lens<RootObject, ChangableObject2>(
      root => {
        const v1 = this.get(root);
        return l2.get(v1 as any);
      },
      (x, newValue) => this.set(x, l2.set(this.get(x) as any, newValue))
    );
  }
}

function lens<RootObject extends object>() {
  function prop<T extends keyof RootObject>(name: T) {
    return Lens.simple<RootObject, T>(name);
  }

  function path<T1 extends keyof RootObject>(
    t1: T1
  ): Lens<RootObject, RootObject[T1]>;
  function path<T1 extends keyof RootObject, T2 extends keyof RootObject[T1]>(
    t1: T1,
    t2: T2
  ): Lens<RootObject, RootObject[T1][T2]>;
  function path<
    T1 extends keyof RootObject,
    T2 extends keyof RootObject[T1],
    T3 extends keyof RootObject[T1][T2]
  >(t1: T1, t2: T2, t3: T3): Lens<RootObject, RootObject[T1][T2][T3]>;
  function path<
    T1 extends keyof RootObject,
    T2 extends keyof RootObject[T1],
    T3 extends keyof RootObject[T1][T2]
  >(t1: T1, t2?: T2, t3?: T3) {
    const l1 = prop(t1);
    if (!t2) return l1;
    const l2 = l1.compose(Lens.simple(t2));
    if (!t3) return l2;
    return l2.compose(Lens.simple(t3));
  }

  return { prop, path };
}

// demo

interface Housing {
  street: string;
  house: number;
  city: string;
}

interface Work {
  name: string;
  location: Housing;
}

interface Employee {
  name: string;
  location: Housing;
  work: Work;
}

const gal: Employee = {
  name: "Gal Schlezinger",
  location: {
    street: "Somewhere",
    house: 0,
    city: "Ramat Gan"
  },
  work: {
    name: "Wix",
    location: {
      street: "Namal Tel Aviv",
      house: 40,
      city: "Tel Aviv"
    }
  }
};

const emp$ = lens<Employee>();
const work$ = lens<Work>();
const housing$ = lens<Housing>();

const emp$work = emp$.prop("work").compose(work$.prop("name"));
console.log(emp$work.over(gal, x => `${x}.com`));

const employeeStreet$ = lens<Employee>().path("location", "street");
console.log(employeeStreet$.get(gal));

const employeeWorkStreet$ = lens<Employee>().path("work", "location", "street");
console.log(employeeWorkStreet$.get(gal));