nix-ip-utils/lib/iterate.nix

{
  lib,
  internal,
  ip,
  cidr,
}:
let
  inherit (builtins) genList;
  inherit (lib) map;
in
rec {
  /**
    Generate a list of IP integers from start to end (inclusive).

    # Type

    ```
    range :: (String | Int) -> (String | Int) -> [Int]
    ```

    # Example

    ```nix
    range "192.168.1.1" "192.168.1.3"
    => [ 3232235777 3232235778 3232235779 ]
    ```
  */
  range =
    start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      count = endInt - startInt + 1;
    in
    if endInt < startInt then [ ] else genList (i: startInt + i) count;

  /**
    Generate a list of IP strings from start to end (inclusive).

    # Type

    ```
    rangeStr :: (String | Int) -> (String | Int) -> [String]
    ```

    # Example

    ```nix
    rangeStr "192.168.1.1" "192.168.1.3"
    => [ "192.168.1.1" "192.168.1.2" "192.168.1.3" ]
    ```
  */
  rangeStr = start: end: map ip.format (range start end);

  /**
    Get all addresses in a CIDR (as integers).

    # Type

    ```
    cidrAddresses :: (String | { address :: Int; prefixLen :: Int; }) -> [Int]
    ```

    # Example

    ```nix
    cidrAddresses "192.168.1.0/30"
    => [ 3232235776 3232235777 3232235778 3232235779 ]
    ```
  */
  cidrAddresses =
    cidrStr:
    let
      net = cidr.network cidrStr;
      bcast = cidr.broadcast cidrStr;
    in
    range net bcast;

  /**
    Get all addresses in a CIDR (as strings).

    # Type

    ```
    cidrAddressesStr :: (String | { address :: Int; prefixLen :: Int; }) -> [String]
    ```

    # Example

    ```nix
    cidrAddressesStr "192.168.1.0/30"
    => [ "192.168.1.0" "192.168.1.1" "192.168.1.2" "192.168.1.3" ]
    ```
  */
  cidrAddressesStr = cidrStr: map ip.format (cidrAddresses cidrStr);

  /**
    Get all usable host addresses in a CIDR (as integers).
    Excludes network and broadcast for /30 and larger.

    # Type

    ```
    cidrHosts :: (String | { address :: Int; prefixLen :: Int; }) -> [Int]
    ```

    # Example

    ```nix
    cidrHosts "192.168.1.0/30"
    => [ 3232235777 3232235778 ]
    ```
  */
  cidrHosts =
    cidrStr:
    let
      c = cidr.parse cidrStr;
      first = cidr.firstHost cidrStr;
      last = cidr.lastHost cidrStr;
    in
    if c.prefixLen >= 31 then range first last else range first last;

  /**
    Get all usable host addresses in a CIDR (as strings).

    # Type

    ```
    cidrHostsStr :: (String | { address :: Int; prefixLen :: Int; }) -> [String]
    ```

    # Example

    ```nix
    cidrHostsStr "192.168.1.0/30"
    => [ "192.168.1.1" "192.168.1.2" ]
    ```
  */
  cidrHostsStr = cidrStr: map ip.format (cidrHosts cidrStr);

  /**
    Check if an IP is within a range (inclusive).

    # Type

    ```
    inRange :: (String | Int) -> (String | Int) -> (String | Int) -> Bool
    ```

    # Example

    ```nix
    inRange "192.168.1.1" "192.168.1.10" "192.168.1.5"
    => true
    ```
  */
  inRange =
    start: end: addr:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      addrInt = ip.parse addr;
    in
    addrInt >= startInt && addrInt <= endInt;

  /**
    Count IPs in a range (inclusive).

    # Type

    ```
    countRange :: (String | Int) -> (String | Int) -> Int
    ```

    # Example

    ```nix
    countRange "192.168.1.1" "192.168.1.10"
    => 10
    ```
  */
  countRange =
    start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
    in
    if endInt < startInt then 0 else endInt - startInt + 1;

  /**
    Map a function over an IP range.
    More memory-efficient than generating the full list first.

    # Type

    ```
    mapRange :: (Int -> a) -> (String | Int) -> (String | Int) -> [a]
    ```

    # Example

    ```nix
    mapRange (i: i * 2) "0.0.0.1" "0.0.0.3"
    => [ 2 4 6 ]
    ```
  */
  mapRange =
    f: start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      count = endInt - startInt + 1;
    in
    if endInt < startInt then [ ] else genList (i: f (startInt + i)) count;

  /**
    Map a function over a CIDR's addresses.

    # Type

    ```
    mapCidr :: (Int -> a) -> (String | { address :: Int; prefixLen :: Int; }) -> [a]
    ```

    # Example

    ```nix
    mapCidr ip.format "192.168.1.0/30"
    => [ "192.168.1.0" "192.168.1.1" "192.168.1.2" "192.168.1.3" ]
    ```
  */
  mapCidr =
    f: cidrStr:
    let
      net = cidr.network cidrStr;
      bcast = cidr.broadcast cidrStr;
    in
    mapRange f net bcast;

  /**
    Map a function over a CIDR's usable hosts.

    # Type

    ```
    mapCidrHosts :: (Int -> a) -> (String | { address :: Int; prefixLen :: Int; }) -> [a]
    ```

    # Example

    ```nix
    mapCidrHosts ip.format "192.168.1.0/30"
    => [ "192.168.1.1" "192.168.1.2" ]
    ```
  */
  mapCidrHosts =
    f: cidrStr:
    let
      first = cidr.firstHost cidrStr;
      last = cidr.lastHost cidrStr;
    in
    mapRange f first last;

  /**
    Filter IPs in a range.

    # Type

    ```
    filterRange :: (Int -> Bool) -> (String | Int) -> (String | Int) -> [Int]
    ```

    # Example

    ```nix
    filterRange (i: lib.mod i 2 == 0) "0.0.0.1" "0.0.0.5"
    => [ 2 4 ]
    ```
  */
  filterRange =
    pred: start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      count = endInt - startInt + 1;
      allIps = if endInt < startInt then [ ] else genList (i: startInt + i) count;
    in
    builtins.filter pred allIps;

  /**
    Find first IP in range matching predicate.

    # Type

    ```
    findInRange :: (Int -> Bool) -> (String | Int) -> (String | Int) -> (Int | Null)
    ```

    # Example

    ```nix
    findInRange (i: lib.mod i 2 == 0) "0.0.0.1" "0.0.0.5"
    => 2
    ```
  */
  findInRange =
    pred: start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      find =
        current:
        if current > endInt then
          null
        else if pred current then
          current
        else
          find (current + 1);
    in
    if endInt < startInt then null else find startInt;

  /**
    Find first IP in range matching predicate (as string).

    # Type

    ```
    findInRangeStr :: (Int -> Bool) -> (String | Int) -> (String | Int) -> (String | Null)
    ```

    # Example

    ```nix
    findInRangeStr (i: lib.mod i 2 == 0) "0.0.0.1" "0.0.0.5"
    => "0.0.0.2"
    ```
  */
  findInRangeStr =
    pred: start: end:
    let
      result = findInRange pred start end;
    in
    if result == null then null else ip.format result;

  /**
    Take first n IPs from a range.

    # Type

    ```
    takeRange :: Int -> (String | Int) -> (String | Int) -> [Int]
    ```

    # Example

    ```nix
    takeRange 2 "192.168.1.1" "192.168.1.10"
    => [ 3232235777 3232235778 ]
    ```
  */
  takeRange =
    n: start: end:
    let
      startInt = ip.parse start;
      endInt = ip.parse end;
      actualEnd = if startInt + n - 1 < endInt then startInt + n - 1 else endInt;
    in
    if endInt < startInt || n <= 0 then [ ] else range startInt actualEnd;

  /**
    Take first n IPs from a range (as strings).

    # Type

    ```
    takeRangeStr :: Int -> (String | Int) -> (String | Int) -> [String]
    ```

    # Example

    ```nix
    takeRangeStr 2 "192.168.1.1" "192.168.1.10"
    => [ "192.168.1.1" "192.168.1.2" ]
    ```
  */
  takeRangeStr =
    n: start: end:
    map ip.format (takeRange n start end);
}