Proposal for a context extension for riverpod, including a working context.watch implementation.

riverpod_context_extension.dart

import 'package:flutter/material.dart';
import 'package:flutter_riverpod/flutter_riverpod.dart';

/*
============
EXAMPLE PART
============
 */

// case 1: basic provider
var cnt1 = StateProvider((ref) => 0, name: 'cnt1');
// case 2: auto dispose
var cnt2 = StateProvider.autoDispose((ref) => 0, name: 'cnt2');
// case 3: family
var cntFamily = Provider.family(
  (ref, int index) => 'family $index ${ref.watch(cnt1)}',
);

void main() {
  runApp(const ProviderScope(
    // Add [InheritedConsumer] underneath the root [ProviderScope]
    child: InheritedConsumer(
      child: MyApp(),
    ),
  ));
}

class MyApp extends StatelessWidget {
  const MyApp({Key? key}) : super(key: key);

  @override
  Widget build(BuildContext context) {
    return MaterialApp(
      title: 'Inherited Consumer Demo',
      home: Scaffold(
        body: Center(
          child: Column(
            mainAxisAlignment: MainAxisAlignment.center,
            children: <Widget>[
              // First counter

              // Wrap in builder to get new context
              // Will only rebuild this builder when cnt1 changes
              Builder(
                builder: (context) => Text('${context.watch(cnt1)}'),
              ),

              // read from anywhere
              IconButton(
                icon: const Icon(Icons.add),
                onPressed: () {
                  context.read(cnt1.state).state++;
                },
              ),
              IconButton(
                icon: const Icon(Icons.remove),
                onPressed: () {
                  context.read(cnt1.state).state--;
                },
              ),

              // Second counter
              Builder(
                builder: (context) {
                  // conditionally display & watch second counter
                  // this won't rebuild when cnt1 > 10 and cnt2 changes
                  if (context.watch(cnt1.select((cnt) => cnt <= 10))) {
                    return Text('${context.watch(cnt2)}');
                  } else {
                    return const Text("Not Watching");
                  }
                },
              ),

              IconButton(
                icon: const Icon(Icons.add),
                onPressed: () {
                  context.read(cnt2.state).state++;
                },
              ),

              Builder(builder: (context) {
                var c1 = context.watch(cnt1);
                var index = c1 <= 10 ? 0 : 1;
                return Text(context.watch(cntFamily(index)));
              })
            ],
          ),
        ),
      ),
    );
  }
}

/*
===================
IMPLEMENTATION PART
===================
 */

extension RiverpodContext on BuildContext {
  /// Nothing special here, taken from the [ConsumerState] implementation
  T read<T>(ProviderBase<T> provider) {
    return ProviderScope.containerOf(this, listen: false).read(provider);
  }

  /// Nothing special here, taken from the [ConsumerState] implementation
  T refresh<T>(ProviderBase<T> provider) {
    return ProviderScope.containerOf(this, listen: false).refresh(provider);
  }

  /// This is new (and special)
  T watch<T>(ProviderListenable<T> provider) {
    return InheritedConsumer.watch<T>(this, provider);
  }
}

class InheritedConsumer extends InheritedWidget {
  const InheritedConsumer({
    required Widget child,
    Key? key,
  }) : super(key: key, child: child);

  @override
  InheritedElement createElement() => _InheritedConsumerElement(this);

  /// This will get the inherited element and call  the
  /// [dependOnInheritedElement] method. The element will decide when to
  /// rebuild the provided context
  static T watch<T>(BuildContext context, ProviderListenable<T> target) {
    var elem =
        context.getElementForInheritedWidgetOfExactType<InheritedConsumer>();
    if (elem == null) {
      throw StateError("No InheritedConsumer found!");
    }
    context.dependOnInheritedElement(elem, aspect: target);
    return (elem as _InheritedConsumerElement)._read(context, target) as T;
  }

  @override
  bool updateShouldNotify(covariant InheritedWidget oldWidget) {
    return true; // This widget should never rebuild, so returning true is fine
  }
}

/// The custom [InheritedElement] implementation. This is somewhat similar to
/// the [ConsumerStatefulElement] implementation.
class _InheritedConsumerElement extends InheritedElement {
  _InheritedConsumerElement(InheritedConsumer widget) : super(widget) {
    // after the initial build we have to set the dependencies
    WidgetsBinding.instance!.addPostFrameCallback(checkDependencies);
    frameCallbackAdded = true;
  }

  /// any active subscriptions on a provider, for each dependent element
  final activeDependents =
      <Element, Map<ProviderListenable, ProviderSubscription>>{};

  /// [InheritedElement] has it's own private [_dependents] but we use both
  /// to identify changes in dependencies after widget rebuilds
  /// (see [checkDependencies])
  final nextDependents = <Element, Set<ProviderListenable>>{};

  /// To keep track of [addPostFrameCallback] calls
  bool frameCallbackAdded = false;

  /// This gets invoked after calling [dependOnInheritedElement] during build.
  /// We use the [aspect] parameter for the provider we want to watch.
  @override
  void updateDependencies(Element dependent, Object? aspect) {
    var listenable = aspect as ProviderListenable;

    if (!activeDependents.containsKey(dependent)) {
      activeDependents[dependent] = {};
    }
    var subscriptions = activeDependents[dependent]!;

    if (!subscriptions.containsKey(listenable)) {
      // create a new [ProviderSubscription] and add it to the dependencies

      void listener(_, v) {
        if (subscriptions[listenable] == null) return;

        // after the build we have to check for changes in dependencies
        WidgetsBinding.instance!.addPostFrameCallback(checkDependencies);
        frameCallbackAdded = true;

        // remove all dependencies, these will be re-assigned
        // during the build phase
        setDependencies(dependent, null);

        // debug print
        print("$listenable CHANGED $v");

        // trigger a rebuild for this dependent
        dependent.markNeedsBuild();
      }

      var container = ProviderScope.containerOf(dependent);
      var subscription = container.listen(listenable, listener);

      // debug print
      print("SUB TO $listenable");

      subscriptions[listenable] = subscription;
    }

    // add the provider to the next dependencies
    (nextDependents[dependent] ??= {}).add(listenable);

    if (!frameCallbackAdded) {
      // add the frame callback if it was not added yet.
      // this happens if the rebuild is not triggered by a provider
      WidgetsBinding.instance!.addPostFrameCallback(checkDependencies);
      frameCallbackAdded = true;
    }
  }

  /// This is used to return the current value when calling [context.watch].
  /// We need this since [dependOnInheritedElement] returns void.
  dynamic _read(Object dependent, ProviderListenable target) {
    return activeDependents[dependent]?[target]?.read();
  }

  @override
  Set<ProviderListenable> getDependencies(Element dependent) {
    return super.getDependencies(dependent) as Set<ProviderListenable>? ?? {};
  }

  /// After building, we have to check all dependencies if they are still
  /// valid and close unused subscriptions
  /// A dependent is automatically removed only from the private dependents
  /// when a widget is deactivated.
  void checkDependencies(Duration _) {
    for (var dependent in activeDependents.entries) {
      // get the next or current dependencies
      var dependencies =
          nextDependents[dependent.key] ?? getDependencies(dependent.key);
      nextDependents.remove(dependent.key);

      // update the current dependencies
      setDependencies(dependent.key, dependencies);

      dependent.value.removeWhere((key, value) {
        // if it was removed during the last build phase, we close
        // the subscription. This is important for auto dispose
        if (!dependencies.contains(key)) {
          // debug print
          print("DESUB FROM $key");

          value.close();
          return true;
        }
        return false;
      });
    }

    activeDependents.removeWhere((key, value) => value.isEmpty);
    frameCallbackAdded = false;
  }

  @override
  void unmount() {
    // cleanup all dependencies
    for (var subscriptions in activeDependents.values) {
      for (var subscription in subscriptions.values) {
        subscription.close();
      }
    }
    super.unmount();
  }
}